respect forced status in datatree simulate for invalid nodes

this fixes a nasty flaky error that I spent many hours tracking down

Author: tybug

hypothesis-python/src/hypothesis/internal/conjecture/data.py

@@ -136,7 +136,8 @@ class BooleanKWargs(TypedDict):
     IntegerKWargs, FloatKWargs, StringKWargs, BytesKWargs, BooleanKWargs
 ]
 IRTypeName: TypeAlias = Literal["integer", "string", "boolean", "float", "bytes"]
-InvalidAt: TypeAlias = Tuple[IRTypeName, IRKWargsType]
+# ir_type, kwargs, forced
+InvalidAt: TypeAlias = Tuple[IRTypeName, IRKWargsType, Optional[IRType]]
 
 
 class ExtraInformation:
@@ -2084,7 +2085,7 @@ def draw_integer(
         )
 
         if self.ir_tree_nodes is not None and observe:
-            node = self._pop_ir_tree_node("integer", kwargs)
+            node = self._pop_ir_tree_node("integer", kwargs, forced=forced)
             if forced is None:
                 assert isinstance(node.value, int)
                 forced = node.value
@@ -2141,7 +2142,7 @@ def draw_float(
         )
 
         if self.ir_tree_nodes is not None and observe:
-            node = self._pop_ir_tree_node("float", kwargs)
+            node = self._pop_ir_tree_node("float", kwargs, forced=forced)
             if forced is None:
                 assert isinstance(node.value, float)
                 forced = node.value
@@ -2183,7 +2184,7 @@ def draw_string(
             },
         )
         if self.ir_tree_nodes is not None and observe:
-            node = self._pop_ir_tree_node("string", kwargs)
+            node = self._pop_ir_tree_node("string", kwargs, forced=forced)
             if forced is None:
                 assert isinstance(node.value, str)
                 forced = node.value
@@ -2219,7 +2220,7 @@ def draw_bytes(
         kwargs: BytesKWargs = self._pooled_kwargs("bytes", {"size": size})
 
         if self.ir_tree_nodes is not None and observe:
-            node = self._pop_ir_tree_node("bytes", kwargs)
+            node = self._pop_ir_tree_node("bytes", kwargs, forced=forced)
             if forced is None:
                 assert isinstance(node.value, bytes)
                 forced = node.value
@@ -2261,7 +2262,7 @@ def draw_boolean(
         kwargs: BooleanKWargs = self._pooled_kwargs("boolean", {"p": p})
 
         if self.ir_tree_nodes is not None and observe:
-            node = self._pop_ir_tree_node("boolean", kwargs)
+            node = self._pop_ir_tree_node("boolean", kwargs, forced=forced)
             if forced is None:
                 assert isinstance(node.value, bool)
                 forced = node.value
@@ -2302,7 +2303,9 @@ def _pooled_kwargs(self, ir_type, kwargs):
             POOLED_KWARGS_CACHE[key] = kwargs
             return kwargs
 
-    def _pop_ir_tree_node(self, ir_type: IRTypeName, kwargs: IRKWargsType) -> IRNode:
+    def _pop_ir_tree_node(
+        self, ir_type: IRTypeName, kwargs: IRKWargsType, *, forced: Optional[IRType]
+    ) -> IRNode:
         assert self.ir_tree_nodes is not None
 
         if self._node_index == len(self.ir_tree_nodes):
@@ -2321,7 +2324,7 @@ def _pop_ir_tree_node(self, ir_type: IRTypeName, kwargs: IRKWargsType) -> IRNode
         # (in fact, it is possible that giving up early here results in more time
         # for useful shrinks to run).
         if node.ir_type != ir_type:
-            invalid_at = (ir_type, kwargs)
+            invalid_at = (ir_type, kwargs, forced)
             self.invalid_at = invalid_at
             self.observer.mark_invalid(invalid_at)
             self.mark_invalid(f"(internal) want a {ir_type} but have a {node.ir_type}")
@@ -2330,7 +2333,7 @@ def _pop_ir_tree_node(self, ir_type: IRTypeName, kwargs: IRKWargsType) -> IRNode
         # that is allowed by the expected kwargs, then we can coerce this node
         # into an aligned one by using its value. It's unclear how useful this is.
         if not ir_value_permitted(node.value, node.ir_type, kwargs):
-            invalid_at = (ir_type, kwargs)
+            invalid_at = (ir_type, kwargs, forced)
             self.invalid_at = invalid_at
             self.observer.mark_invalid(invalid_at)
             self.mark_invalid(f"(internal) got a {ir_type} but outside the valid range")

hypothesis-python/src/hypothesis/internal/conjecture/datatree.py

@@ -840,8 +840,8 @@ def simulate_test_function(self, data):
         tree. This will likely change in future."""
         node = self.root
 
-        def draw(ir_type, kwargs, *, forced=None):
-            if ir_type == "float" and forced is not None:
+        def draw(ir_type, kwargs, *, forced=None, convert_forced=True):
+            if ir_type == "float" and forced is not None and convert_forced:
                 forced = int_to_float(forced)
 
             draw_func = getattr(data, f"draw_{ir_type}")
@@ -866,9 +866,9 @@ def draw(ir_type, kwargs, *, forced=None):
                     data.conclude_test(t.status, t.interesting_origin)
                 elif node.transition is None:
                     if node.invalid_at is not None:
-                        (ir_type, kwargs) = node.invalid_at
+                        (ir_type, kwargs, forced) = node.invalid_at
                         try:
-                            draw(ir_type, kwargs)
+                            draw(ir_type, kwargs, forced=forced, convert_forced=False)
                         except StopTest:
                             if data.invalid_at is not None:
                                 raise

hypothesis-python/tests/conjecture/test_data_tree.py

@@ -621,8 +621,8 @@ def test_datatree_repr(bool_kwargs, int_kwargs):
     )
 
 
-def _draw(data, node):
-    return getattr(data, f"draw_{node.ir_type}")(**node.kwargs)
+def _draw(data, node, *, forced=None):
+    return getattr(data, f"draw_{node.ir_type}")(**node.kwargs, forced=forced)
 
 
 @given(ir_nodes(), ir_nodes())
@@ -641,7 +641,7 @@ def test_misaligned_nodes_after_valid_draw(node, misaligned_node):
     tree.simulate_test_function(data)
     assert data.status is Status.INVALID
 
-    assert data.invalid_at == (node.ir_type, node.kwargs)
+    assert data.invalid_at == (node.ir_type, node.kwargs, None)
 
 
 @given(ir_nodes(was_forced=False), ir_nodes(was_forced=False))
@@ -709,3 +709,37 @@ def test_simulate_non_invalid_conclude_is_unseen_behavior(node, misaligned_node)
         tree.simulate_test_function(data)
 
     assert data.status is Status.OVERRUN
+
+
+@given(ir_nodes(), ir_nodes())
+@settings(suppress_health_check=[HealthCheck.too_slow])
+def test_simulating_inherits_invalid_forced_status(node, misaligned_node):
+    assume(misaligned_node.ir_type != node.ir_type)
+
+    # we have some logic in DataTree.simulate_test_function to "peak ahead" and
+    # make sure it simulates invalid nodes correctly. But if it does so without
+    # respecting whether the invalid node was forced or not, and this simulation
+    # is observed by an observer, this can cause flaky errors later due to a node
+    # going from unforced to forced.
+
+    tree = DataTree()
+
+    def test_function(ir_nodes):
+        data = ConjectureData.for_ir_tree(ir_nodes, observer=tree.new_observer())
+        _draw(data, node)
+        _draw(data, node, forced=node.value)
+
+    # (1) set up a misaligned node at index 1
+    with pytest.raises(StopTest):
+        test_function([node, misaligned_node])
+
+    # (2) simulate an aligned tree. the datatree peaks ahead here using invalid_at
+    # due to (1).
+    data = ConjectureData.for_ir_tree([node, node], observer=tree.new_observer())
+    with pytest.raises(PreviouslyUnseenBehaviour):
+        tree.simulate_test_function(data)
+
+    # (3) run the same aligned tree without simulating. this uses the actual test
+    # function's draw and forced value. This would flaky error if it did not match
+    # what the datatree peaked ahead with in (2).
+    test_function([node, node])