Add assert always and eventually (#22)

* Add assert always and eventually

* Pump version

Author: Tomperez98

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "resonate-sdk"
-version = "0.1.13"
+version = "0.1.14"
 description = "Distributed Async Await by Resonate HQ, Inc"
 authors = [
     { name = "Resonate HQ, Inc", email = "contact@resonatehq.io" }

src/resonate/dst/scheduler.py

@@ -115,7 +115,11 @@ def __init__(  # noqa: PLR0913
         failure_chance: float,
         mode: Mode,
         probe: Callable[[Dependencies, int], Any] | None,
+        assert_eventually: Callable[[Dependencies, int], None] | None,
+        assert_always: Callable[[Dependencies, int], None] | None,
     ) -> None:
+        self._assert_eventually = assert_eventually
+        self._assert_always = assert_always
         self._runnable_coroutines: RunnableCoroutines = []
         self._awaitables: Awaitables = {}
         self._runnable_functions: RunnableFunctions = []
@@ -333,7 +337,7 @@ def _run_function_and_move_awaitables_to_runnables(
             v=v,
         )
 
-    def _run(self) -> list[Promise[Any]]:
+    def _run(self) -> list[Promise[Any]]:  # noqa: C901
         promises = self._initialize_runnables()
 
         while True:
@@ -347,6 +351,10 @@ def _run(self) -> list[Promise[Any]]:
 
             if self._probe is not None:
                 self._probe_results.append(self._probe(self.deps, self.tick))
+
+            if self._assert_always is not None:
+                self._assert_always(self.deps, self.tick)
+
             self.tick += 1
 
             if (
@@ -370,11 +378,19 @@ def _run(self) -> list[Promise[Any]]:
             else:
                 assert_never(next_step)
 
+        if self._assert_eventually is not None:
+            self._assert_eventually(self.deps, self.tick)
+
         assert all(p.done() for p in promises), "All promises should be resolved."
         if self._log_file is not None:
             self.dump(file=self._log_file)
         return promises
 
+    def _add_coro_to_awaitables(
+        self, p: Promise[Any], coro_and_promise: CoroAndPromise[Any]
+    ) -> None:
+        self._awaitables.setdefault(p, []).append(coro_and_promise)
+
     def _process_each_runnable(
         self,
         runnable: Runnable[Any],
@@ -414,7 +430,9 @@ def _process_each_runnable(
                 invocation=yieldable_or_final_value.to_invoke(),
                 runnable=runnable,
             )
-            self._awaitables[p] = [runnable.coro_and_promise]
+            self._add_coro_to_awaitables(
+                p=p, coro_and_promise=runnable.coro_and_promise
+            )
             self._events.append(
                 AwaitedForPromise(promise_id=p.promise_id, tick=self.tick)
             )
@@ -427,8 +445,8 @@ def _process_each_runnable(
             self._runnable_coroutines.append(Runnable(runnable.coro_and_promise, Ok(p)))
 
         elif isinstance(yieldable_or_final_value, Promise):
-            self._awaitables.setdefault(yieldable_or_final_value, []).append(
-                runnable.coro_and_promise,
+            self._add_coro_to_awaitables(
+                yieldable_or_final_value, runnable.coro_and_promise
             )
             if yieldable_or_final_value.done():
                 unblock_depands_coros(

src/resonate/testing.py

@@ -27,6 +27,8 @@ def dst(  # noqa: PLR0913
     failure_chance: float = 0,
     max_failures: int = 0,
     probe: Callable[[Dependencies, int], Any] | None = None,
+    assert_always: Callable[[Dependencies, int], Any] | None = None,
+    assert_eventually: Callable[[Dependencies, int], Any] | None = None,
 ) -> list[DSTScheduler]:
     schedulers: list[DSTScheduler] = []
 
@@ -45,6 +47,8 @@ def dst(  # noqa: PLR0913
                     max_failures=max_failures,
                     log_file=log_file,
                     probe=probe,
+                    assert_always=assert_always,
+                    assert_eventually=assert_eventually,
                 )
                 for i in seed
             )
@@ -58,6 +62,8 @@ def dst(  # noqa: PLR0913
                     max_failures=max_failures,
                     log_file=log_file,
                     probe=probe,
+                    assert_always=assert_always,
+                    assert_eventually=assert_eventually,
                 )
             )
         else: