From 7cafa11a83ca04caf4c1cfbf5cd0a091c5780f28 Mon Sep 17 00:00:00 2001
From: Andras Schmelczer <andras@schmelczer.dev>
Date: Sat, 6 Jun 2026 21:57:27 +0100
Subject: [PATCH] Fix CI

---
 great_ai/deploy/great_ai.py      | 18 +++++++++---------
 tests/external/test_internals.py | 26 +++++++++++++++-----------
 2 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/great_ai/deploy/great_ai.py b/great_ai/deploy/great_ai.py
index 55b2125..a6ea818 100644
--- a/great_ai/deploy/great_ai.py
+++ b/great_ai/deploy/great_ai.py
@@ -230,17 +230,17 @@ class GreatAI(Generic[T, V]):
                 ),
             )
 
+        # inner/inner_async return the class' T (bound to Trace | Awaitable[Trace]);
+        # in this method T resolves to Trace[V], but that is not provable to mypy, so
+        # cast to the concrete shape parallel_map expects.
+        map_function = cast(
+            Callable[[Any], Union[Trace[V], Awaitable[Trace[V]]]],
+            inner_async if get_function_metadata_store(self).is_asynchronous else inner,
+        )
+
         return list(
             tqdm(
-                parallel_map(
-                    (
-                        inner_async
-                        if get_function_metadata_store(self).is_asynchronous
-                        else inner
-                    ),
-                    batch,
-                    concurrency=concurrency,
-                ),
+                parallel_map(map_function, batch, concurrency=concurrency),
                 total=len(batch),
             )
         )
diff --git a/tests/external/test_internals.py b/tests/external/test_internals.py
index d858c17..142c1fa 100644
--- a/tests/external/test_internals.py
+++ b/tests/external/test_internals.py
@@ -242,7 +242,7 @@ async def test_wait_closed():
             return_exceptions=True,
         )
         assert ret == []
-        assert mocked.called_once()
+        mocked.assert_called_once()
 
     asyncio.set_event_loop(loop)
     with mock.patch("great_ai.external.async_lru._close_waited") as mocked:
@@ -251,7 +251,7 @@ async def test_wait_closed():
             return_exceptions=True,
         )
         assert ret == []
-        assert mocked.called_once()
+        mocked.assert_called_once()
     asyncio.set_event_loop(None)
 
     fut = loop.create_future()
@@ -263,7 +263,7 @@ async def test_wait_closed():
             return_exceptions=True,
         )
         assert ret == [None]
-        assert mocked.called_once()
+        mocked.assert_called_once()
 
     exc = ZeroDivisionError()
     fut = loop.create_future()
@@ -275,7 +275,7 @@ async def test_wait_closed():
             return_exceptions=True,
         )
         assert ret == [exc]
-        assert mocked.called_once()
+        mocked.assert_called_once()
 
     fut = loop.create_future()
     fut.set_exception(ZeroDivisionError)
@@ -286,17 +286,21 @@ async def test_wait_closed():
                 wrapped,
                 return_exceptions=False,
             )
-        assert mocked.called_once()
+        # the awaited result raised before _wait_closed could flush its done
+        # callback, so yield once to let the scheduled _close_waited run
+        await asyncio.sleep(0)
+        mocked.assert_called_once()
 
 
 def test_close_waited():
     wrapped = Wrapped()
-    wrapped.cache_clear = partial(_cache_clear, wrapped)
+    # _close_waited calls wrapped.cache_clear(); mock that directly (patching the
+    # module-level _cache_clear would not affect the already-bound partial).
+    wrapped.cache_clear = mock.Mock()
 
-    with mock.patch("great_ai.external.async_lru._cache_clear") as mocked:
-        _close_waited(wrapped, None)
+    _close_waited(wrapped, None)
 
-        assert mocked.called_once()
+    wrapped.cache_clear.assert_called_once()
 
 
 def test_cache_info():
@@ -343,7 +347,7 @@ def test_cache_hit():
     with mock.patch("great_ai.external.async_lru.__cache_touch") as mocked:
         _cache_hit(wrapped, 1)
 
-        assert mocked.called_once()
+        mocked.assert_called_once()
 
     assert wrapped.hits == 2
 
@@ -361,7 +365,7 @@ def test_cache_miss():
     with mock.patch("great_ai.external.async_lru.__cache_touch") as mocked:
         _cache_miss(wrapped, 1)
 
-        assert mocked.called_once()
+        mocked.assert_called_once()
 
     assert wrapped.misses == 2