From e76f49c7f0eb3b9cff9d5641384a441b5a121c67 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Sun, 22 Sep 2024 12:37:24 -0700
Subject: [PATCH] Fix flaky test

---
 libs/langgraph/tests/any_str.py           | 17 +++++++++++++++++
 libs/langgraph/tests/test_pregel.py       | 12 ++++++------
 libs/langgraph/tests/test_pregel_async.py | 12 ++++++------
 3 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/libs/langgraph/tests/any_str.py b/libs/langgraph/tests/any_str.py
index 9a1977a8c..5995d0e52 100644
--- a/libs/langgraph/tests/any_str.py
+++ b/libs/langgraph/tests/any_str.py
@@ -2,6 +2,23 @@ import re
 from typing import Any, Sequence, Union
 
 
+class FloatBetween(float):
+    def __init__(self, min_value: float, max_value: float) -> None:
+        super().__init__()
+        self.min_value = min_value
+        self.max_value = max_value
+
+    def __eq__(self, other: object) -> bool:
+        return (
+            isinstance(other, float)
+            and other >= self.min_value
+            and other <= self.max_value
+        )
+
+    def __hash__(self) -> int:
+        return hash((float(self), self.min_value, self.max_value))
+
+
 class AnyStr(str):
     def __init__(self, prefix: Union[str, re.Pattern] = "") -> None:
         super().__init__()
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 18bccbe75..605f00747 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -72,7 +72,7 @@ from langgraph.pregel import (
 from langgraph.pregel.retry import RetryPolicy
 from langgraph.store.memory import MemoryStore
 from langgraph.types import Interrupt, PregelTask, Send, StreamWriter
-from tests.any_str import AnyDict, AnyStr, AnyVersion, UnsortedSequence
+from tests.any_str import AnyDict, AnyStr, AnyVersion, FloatBetween, UnsortedSequence
 from tests.conftest import ALL_CHECKPOINTERS_SYNC, SHOULD_CHECK_SNAPSHOTS
 from tests.fake_chat import FakeChatModel
 from tests.fake_tracer import FakeTracer
@@ -8596,22 +8596,22 @@ def test_stream_subgraphs_during_execution(
     assert chunks == [
         # arrives before "inner" finishes
         (
-            0.0,
+            FloatBetween(0.0, 0.1),
             (
                 (AnyStr("inner:"),),
                 {"inner_1": {"my_key": "got here", "my_other_key": ""}},
             ),
         ),
-        (0.2, ((), {"outer_1": {"my_key": " and parallel"}})),
+        (FloatBetween(0.2, 0.3), ((), {"outer_1": {"my_key": " and parallel"}})),
         (
-            0.5,
+            FloatBetween(0.5, 0.6),
             (
                 (AnyStr("inner:"),),
                 {"inner_2": {"my_key": " and there", "my_other_key": "got here"}},
             ),
         ),
-        (0.5, ((), {"inner": {"my_key": "got here and there"}})),
-        (0.5, ((), {"outer_2": {"my_key": " and back again"}})),
+        (FloatBetween(0.5, 0.6), ((), {"inner": {"my_key": "got here and there"}})),
+        (FloatBetween(0.5, 0.6), ((), {"outer_2": {"my_key": " and back again"}})),
     ]
 
 
diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
index 0bd9ed1f9..d17925aca 100644
--- a/libs/langgraph/tests/test_pregel_async.py
+++ b/libs/langgraph/tests/test_pregel_async.py
@@ -70,7 +70,7 @@ from langgraph.pregel import (
 from langgraph.pregel.retry import RetryPolicy
 from langgraph.store.memory import MemoryStore
 from langgraph.types import Interrupt, PregelTask, Send, StreamWriter
-from tests.any_str import AnyDict, AnyStr, AnyVersion, UnsortedSequence
+from tests.any_str import AnyDict, AnyStr, AnyVersion, FloatBetween, UnsortedSequence
 from tests.conftest import (
     ALL_CHECKPOINTERS_ASYNC,
     ALL_CHECKPOINTERS_ASYNC_PLUS_NONE,
@@ -7202,22 +7202,22 @@ async def test_stream_subgraphs_during_execution(checkpointer_name: str) -> None
         assert chunks == [
             # arrives before "inner" finishes
             (
-                0.0,
+                FloatBetween(0.0, 0.1),
                 (
                     (AnyStr("inner:"),),
                     {"inner_1": {"my_key": "got here", "my_other_key": ""}},
                 ),
             ),
-            (0.2, ((), {"outer_1": {"my_key": " and parallel"}})),
+            (FloatBetween(0.2, 0.3), ((), {"outer_1": {"my_key": " and parallel"}})),
             (
-                0.5,
+                FloatBetween(0.5, 0.6),
                 (
                     (AnyStr("inner:"),),
                     {"inner_2": {"my_key": " and there", "my_other_key": "got here"}},
                 ),
             ),
-            (0.5, ((), {"inner": {"my_key": "got here and there"}})),
-            (0.5, ((), {"outer_2": {"my_key": " and back again"}})),
+            (FloatBetween(0.5, 0.6), ((), {"inner": {"my_key": "got here and there"}})),
+            (FloatBetween(0.5, 0.6), ((), {"outer_2": {"my_key": " and back again"}})),
         ]