Throw exceptions when inserting duplicate streams.

Lukasa · Lukasa · commit 3ede14983798 · 2016-01-06T16:44:18.000Z
diff --git a/docs/source/api.rst b/docs/source/api.rst
@@ -13,3 +13,5 @@ Exceptions
 .. autoclass:: priority.DeadlockError
 
 .. autoclass:: priority.PriorityLoop
+
+.. autoclass:: priority.DuplicateStreamError
diff --git a/src/priority/__init__.py b/src/priority/__init__.py
@@ -2,4 +2,6 @@
 """
 priority: HTTP/2 priority implementation for Python
 """
-from .priority import Stream, PriorityTree, DeadlockError, PriorityLoop  # noqa
+from .priority import (  # noqa
+    Stream, PriorityTree, DeadlockError, PriorityLoop, DuplicateStreamError
+)
diff --git a/src/priority/priority.py b/src/priority/priority.py
@@ -33,6 +33,13 @@ class PriorityLoop(Exception):
     pass
 
 
+class DuplicateStreamError(Exception):
+    """
+    An attempt was made to insert a stream that already exists.
+    """
+    pass
+
+
 class Stream(object):
     """
     Priority information for a given stream.
@@ -220,7 +227,9 @@ def insert_stream(self,
         :param exclusive: (optional) Whether this new stream should be an
             exclusive dependency of the parent.
         """
-        assert stream_id not in self._streams
+        if stream_id in self._streams:
+            raise DuplicateStreamError("Stream %d already in tree", stream_id)
+
         stream = Stream(stream_id, weight)
 
         if exclusive:
diff --git a/test/test_priority.py b/test/test_priority.py
@@ -262,6 +262,17 @@ def test_can_reprioritize_a_stream(self,
         actual_result = [next(t) for _ in range(len(result))]
         assert actual_result == result
 
+    def test_priority_tree_raises_error_inserting_duplicate(self):
+        """
+        Attempting to insert a stream that is already in the tree raises a
+        DuplicateStreamError
+        """
+        p = priority.PriorityTree()
+        p.insert_stream(1)
+
+        with pytest.raises(priority.DuplicateStreamError):
+            p.insert_stream(1)
+
 
 class TestPriorityTreeOutput(object):
     """
