Add MissingStreamError exception

Lukasa · Lukasa · commit adc1720bb59f · 2016-01-06T16:56:06.000Z
diff --git a/docs/source/api.rst b/docs/source/api.rst
@@ -15,3 +15,5 @@ Exceptions
 .. autoclass:: priority.PriorityLoop
 
 .. autoclass:: priority.DuplicateStreamError
+
+.. autoclass:: priority.MissingStreamError
diff --git a/src/priority/__init__.py b/src/priority/__init__.py
@@ -3,5 +3,6 @@
 priority: HTTP/2 priority implementation for Python
 """
 from .priority import (  # noqa
-    Stream, PriorityTree, DeadlockError, PriorityLoop, DuplicateStreamError
+    Stream, PriorityTree, DeadlockError, PriorityLoop, DuplicateStreamError,
+    MissingStreamError
 )
diff --git a/src/priority/priority.py b/src/priority/priority.py
@@ -40,6 +40,13 @@ class DuplicateStreamError(Exception):
     pass
 
 
+class MissingStreamError(KeyError, Exception):
+    """
+    An operation was attempted on a stream that is not present in the tree.
+    """
+    pass
+
+
 class Stream(object):
     """
     Priority information for a given stream.
@@ -282,7 +289,10 @@ def stream_cycle(new_parent, current):
                 "Stream %d is in a priority loop." % new_parent.stream_id
             )  # pragma: no cover
 
-        current_stream = self._streams[stream_id]
+        try:
+            current_stream = self._streams[stream_id]
+        except KeyError:
+            raise MissingStreamError("Stream %d not in tree" % stream_id)
 
         # Update things in a specific order to make sure the calculation
         # behaves properly. Specifically, we first update the weight. Then,
@@ -321,7 +331,11 @@ def remove_stream(self, stream_id):
 
         :param stream_id: The ID of the stream to remove.
         """
-        child = self._streams.pop(stream_id)
+        try:
+            child = self._streams.pop(stream_id)
+        except KeyError:
+            raise MissingStreamError("Stream %d not in tree" % stream_id)
+
         parent = child.parent
         parent.remove_child(child)
 
@@ -331,7 +345,10 @@ def block(self, stream_id):
 
         :param stream_id: The ID of the stream to block.
         """
-        self._streams[stream_id].active = False
+        try:
+            self._streams[stream_id].active = False
+        except KeyError:
+            raise MissingStreamError("Stream %d not in tree" % stream_id)
 
     def unblock(self, stream_id):
         """
@@ -340,7 +357,10 @@ def unblock(self, stream_id):
         :param stream_id: The ID of the stream to unblock.
         """
         # When a stream becomes unblocked,
-        self._streams[stream_id].active = True
+        try:
+            self._streams[stream_id].active = True
+        except KeyError:
+            raise MissingStreamError("Stream %d not in tree" % stream_id)
 
     # The iterator protocol
     def __iter__(self):  # pragma: no cover
diff --git a/test/test_priority.py b/test/test_priority.py
@@ -273,6 +273,25 @@ def test_priority_tree_raises_error_inserting_duplicate(self):
         with pytest.raises(priority.DuplicateStreamError):
             p.insert_stream(1)
 
+    def test_priority_raises_good_errors_for_missing_streams(self):
+        """
+        Attempting operations on absent streams raises a MissingStreamError.
+        """
+        p = priority.PriorityTree()
+        p.insert_stream(1)
+
+        with pytest.raises(priority.MissingStreamError):
+            p.reprioritize(3)
+
+        with pytest.raises(priority.MissingStreamError):
+            p.block(3)
+
+        with pytest.raises(priority.MissingStreamError):
+            p.unblock(3)
+
+        with pytest.raises(priority.MissingStreamError):
+            p.remove_stream(3)
+
 
 class TestPriorityTreeOutput(object):
     """

Original file line number	Diff line number	Diff line change
`@@ -3,5 +3,6 @@`
`3`	`3`	`priority: HTTP/2 priority implementation for Python`
`4`	`4`	`"""`
`5`	`5`	`from .priority import ( # noqa`
`6`		`- Stream, PriorityTree, DeadlockError, PriorityLoop, DuplicateStreamError`
	`6`	`+ Stream, PriorityTree, DeadlockError, PriorityLoop, DuplicateStreamError,`
	`7`	`+ MissingStreamError`
`7`	`8`	`)`