chore(py): fix integration tests (#1543)

python/tests/integration_tests/test_async_client.py

@@ -48,11 +48,11 @@ async def check_similar_examples():
 
 
 # Helper function to wait for a condition
-async def wait_for(condition, timeout=10):
+async def wait_for(condition, timeout=10, **kwargs):
     start_time = asyncio.get_event_loop().time()
     while True:
         try:
-            if await condition():
+            if await condition(**kwargs):
                 return
         except Exception:
             if asyncio.get_event_loop().time() - start_time > timeout:
@@ -415,16 +415,24 @@ async def test_annotation_queue_runs(async_client: AsyncClient):
 
     # Create some test runs
     run_ids = [uuid.uuid4() for _ in range(3)]
-    for i in range(3):
+    for i, run_id in enumerate(run_ids):
         await async_client.create_run(
             name=f"test_run_{i}",
             inputs={"input": f"test_{i}"},
             run_type="llm",
             project_name=project_name,
             start_time=datetime.datetime.now(datetime.timezone.utc),
-            id=run_ids[i],
+            id=run_id,
         )
 
+    async def _get_run(run_id: uuid.UUID) -> bool:
+        try:
+            await async_client.read_run(run_id)  # type: ignore
+            return True
+        except ls_utils.LangSmithError:
+            return False
+
+    await asyncio.gather(*[wait_for(_get_run, run_id=run_id) for run_id in run_ids])
     # Add runs to queue
     await async_client.add_runs_to_annotation_queue(queue_id=queue.id, run_ids=run_ids)