tester: Process message contexts

src/saturn_engine/utils/tester/inventory_test.py

@@ -32,10 +32,11 @@ def run_saturn_inventory(
     async def run_inventory() -> None:
         count = 0
         async for item in inventory.iterate(after=after):
-            items.append(asdict(item.as_topic_message()))
-            count = count + 1
-            if limit and count >= limit:
-                break
+            async with item:
+                items.append(asdict(item.as_topic_message()))
+                count = count + 1
+                if limit and count >= limit:
+                    break
 
     asyncio.run(run_inventory())
 

src/saturn_engine/utils/tester/topic_test.py

@@ -1,4 +1,3 @@
-from typing import AsyncContextManager
 from typing import Optional
 
 import asyncio
@@ -33,17 +32,15 @@ async def run_topic() -> None:
         count = 0
         skipped = 0
         async for message in topic.run():
-            if skip is not None and skipped < skip:
-                skipped += 1
-                continue
+            async with message:
+                if skip is not None and skipped < skip:
+                    skipped += 1
+                    continue
 
-            if isinstance(message, AsyncContextManager):
-                message = await message.__aenter__()
-
-            messages.append(asdict(message))
-            count = count + 1
-            if limit and count >= limit:
-                break
+                messages.append(asdict(message))
+                count = count + 1
+                if limit and count >= limit:
+                    break
 
     asyncio.run(run_topic())