python-hyper
diff --git a/‎src/priority/priority.py‎
Lines changed: 6 additions & 3 deletions b/‎src/priority/priority.py‎
Lines changed: 6 additions & 3 deletions
diff --git a/‎test/test_priority.py‎
Lines changed: 50 additions & 1 deletion b/‎test/test_priority.py‎
Lines changed: 50 additions & 1 deletion
@@ -82,6 +82,9 @@ def remove_child(self, child):
 
         self.child_queue = new_queue
 
+        for new_child in child.children:
+            self.add_child(new_child)
+
     def schedule(self):
         """
         Returns the stream ID of the next child to schedule. Potentially
@@ -217,9 +220,9 @@ def remove_stream(self, stream_id):
         """
         Removes a stream from the priority tree.
         """
-        # TODO: At some point we should actually prune streams we no longer
-        # need. For now, just mark it permanently blocked.
-        self._streams[stream_id].active = False
+        child = self._streams.pop(stream_id)
+        parent = child.parent
+        parent.remove_child(child)
 
     def block(self, stream_id):
         """
 
@@ -44,7 +44,6 @@
 )
 
 
-
 def readme_tree():
     """
     Provide a tree configured as the one in the readme.
@@ -113,6 +112,38 @@ def test_stream_repr(self):
         s = priority.Stream(stream_id=80, weight=16)
         assert repr(s) == "Stream<id=80, weight=16>"
 
+    @given(STREAMS_AND_WEIGHTS)
+    def test_streams_are_well_ordered(self, streams_and_weights):
+        """
+        Streams are ordered by their stream ID.
+        """
+        stream_list = [
+            priority.Stream(stream_id=s, weight=w)
+            for s, w in streams_and_weights
+        ]
+        stream_list = sorted(stream_list)
+        streams_by_id = [stream.stream_id for stream in stream_list]
+        assert sorted(streams_by_id) == streams_by_id
+
+    @given(
+        integers(min_value=1, max_value=2**24),
+        integers(min_value=1, max_value=2**24)
+    )
+    def test_stream_ordering(self, a, b):
+        """
+        Two streams are well ordered based on their stream ID.
+        """
+        s1 = priority.Stream(stream_id=a, weight=16)
+        s2 = priority.Stream(stream_id=b, weight=32)
+
+        assert (s1 < s2) == (a < b)
+        assert (s1 <= s2) == (a <= b)
+        assert (s1 > s2) == (a > b)
+        assert (s1 >= s2) == (a >= b)
+        assert (s1 == s2) == (a == b)
+        assert (s1 != s2) == (a != b)
+
+
 
 class TestPriorityTreeManual(object):
     """
@@ -161,6 +192,24 @@ def test_priority_tree_blocking_is_isomorphic(self,
         result = [next(tree) for _ in range(len(expected))]
         assert expected == result
 
+    @given(BLOCKED_AND_ACTIVE)
+    def test_removing_items_behaves_similarly_to_blocking(self,
+                                                          blocked_expected):
+        """
+        From the perspective of iterating over items, removing streams should
+        have the same effect as blocking them, except that the ordering
+        changes. Because the ordering is not important, don't test for it.
+        """
+        tree = readme_tree()
+        blocked = blocked_expected[0]
+        expected = set(blocked_expected[1])
+
+        for stream_id in blocked:
+            tree.remove_stream(stream_id)
+
+        result = set(next(tree) for _ in range(len(expected)))
+        assert expected == result
+
     def test_priority_tree_raises_deadlock_error_if_all_blocked(self):
         """
         Assuming all streams are blocked and none can progress, asking for the