From 38098e70c16a8e851ae063b4583e0e369f1767a4 Mon Sep 17 00:00:00 2001
From: Simon Alibert <simon.alibert@huggingface.co>
Date: Sat, 22 Mar 2025 13:45:04 +0100
Subject: [PATCH 1/3] Add wait_calls & wait_called

---
 src/mock_serial/stub.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/src/mock_serial/stub.py b/src/mock_serial/stub.py
index 6e98018..addcfa3 100644
--- a/src/mock_serial/stub.py
+++ b/src/mock_serial/stub.py
@@ -1,3 +1,7 @@
+import threading
+import time
+
+
 class Stub:
     def __init__(
         self,
@@ -9,6 +13,7 @@ def __init__(
         self.__receive_bytes = receive_bytes
         self.__send_bytes = send_bytes
         self.__calls = 0
+        self._event = threading.Event()
 
         if (send_bytes is None) == (send_fn is None):
             raise TypeError('Specify either send_bytes or send_fn.')
@@ -27,16 +32,32 @@ def receive_bytes(self):
 
     def call(self):
         self.__calls += 1
+        self._event.set()
         return self.__send_fn(self.__calls)
 
     @property
     def calls(self):
         return self.__calls
 
+    def wait_calls(
+        self, min_calls: int = 1, timeout: float = 1.0, sleeps: float = 0.005
+    ):
+        start = time.perf_counter()
+        while time.perf_counter() - start < timeout:
+            if self.calls >= min_calls:
+                return self.calls
+            time.sleep(sleeps)
+        raise TimeoutError(
+            f"Stub not called {min_calls} times within {timeout} seconds."
+        )
+
     @property
     def called(self):
         return self.calls > 0
 
+    def wait_called(self, timeout: float = 1.0):
+        return self._event.wait(timeout)
+
     def __repr__(self):
         if self.__send_bytes:
             return f"{self.receive_bytes} => {self.__send_bytes}"

From 3c33b1cce88c72a23dcba5ba8720ab99342b5d74 Mon Sep 17 00:00:00 2001
From: Simon Alibert <simon.alibert@huggingface.co>
Date: Sat, 22 Mar 2025 13:46:23 +0100
Subject: [PATCH 2/3] Add tests

---
 tests/test_stub.py | 44 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/tests/test_stub.py b/tests/test_stub.py
index a120ab1..0c2a233 100644
--- a/tests/test_stub.py
+++ b/tests/test_stub.py
@@ -1,3 +1,6 @@
+import threading
+import time
+
 import pytest
 
 from mock_serial.stub import Stub
@@ -23,3 +26,44 @@ def test_init_bad_args(kwargs):
 
     msg = 'Specify either send_bytes or send_fn.'
     assert str(excinfo.value) == msg
+
+
+def test_wait_called_success():
+    stub = Stub(receive_bytes=b'123', send_bytes=b'456')
+
+    def trigger():
+        time.sleep(0.01)
+        stub.call()
+
+    threading.Thread(target=trigger).start()
+    assert stub.wait_called(timeout=0.1)
+
+
+def test_wait_called_timeout():
+    stub = Stub(receive_bytes=b'123', send_bytes=b'456')
+    assert stub.wait_called(timeout=0.05) is False
+
+
+def test_wait_calls_success():
+    stub = Stub(receive_bytes=b'123', send_bytes=b'456')
+
+    def trigger():
+        for _ in range(3):
+            time.sleep(0.01)
+            stub.call()
+
+    threading.Thread(target=trigger).start()
+    assert stub.wait_calls(min_calls=3, timeout=0.2) == 3
+
+
+def test_wait_calls_timeout():
+    stub = Stub(receive_bytes=b'123', send_bytes=b'456')
+
+    def trigger():
+        for _ in range(2):
+            time.sleep(0.01)
+            stub.call()
+
+    threading.Thread(target=trigger).start()
+    with pytest.raises(TimeoutError, match=r"Stub not called 3 times.*"):
+        stub.wait_calls(min_calls=3, timeout=0.1)

From 041819453d9b4f1885c2a5405e31e69aa6505fbf Mon Sep 17 00:00:00 2001
From: Simon Alibert <simon.alibert@huggingface.co>
Date: Sat, 22 Mar 2025 14:01:35 +0100
Subject: [PATCH 3/3] Add doc

---
 README.md | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/README.md b/README.md
index 390845a..4cdf62f 100644
--- a/README.md
+++ b/README.md
@@ -86,6 +86,25 @@ assert device.stubs['foo'].called
 
 Note: you must ensure there is a stub to match each part of the byte sequence, **otherwise MockSerial will stop responding**. MockSerial does not support "regex" or "placeholder" matching.
 
+In some situations, a test might be checking if a stub has been called *before* the MockSerial thread had time to read, match, and call the stub.
+To prevent this, you can use `wait_called()` and `wait_calls()` which are the waitable equivalent of `called` and `calls`.
+
+```python
+stub = device.stub(
+  receive_bytes=b'123',
+  send_bytes=b'456'
+)
+
+...
+
+# wait until the stub is called at least once
+assert device.stubs['foo'].wait_called(timeout=1.0)
+
+# or wait until it's called a specific number of times
+assert device.stubs['foo'].wait_calls(min_calls=3, timeout=1.0) == 3
+```
+> Note: Contrary to `called` and `calls`, these are methods and not properties so they must be called with parenthesis `()`.
+
 ## Advanced
 
 MockSerial supports overriding stubs by `name` or `receive_bytes`. This can be useful if you want to define most of your stubs once, but override the `send_bytes` for one or two of them in specific tests.