fix: resolve flaky tests across platforms and Python versions, appearing on main branch (#649)

* fix: update test parameter for collate_requests to improve test reliability

* fix: increase wait time in test_worker_restart_and_server_shutdown for macOS compatibility

* fix: update concurrency cancel-in-progress condition for pull requests

* test with skip test

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add timeout to these fn

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update test

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: bhimrazy

.github/workflows/ci-testing.yml

@@ -2,7 +2,7 @@ name: CI testing
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 on:
   push:

tests/unit/test_batch.py

@@ -248,7 +248,7 @@ def test_batched_loop():
         pytest.param(0, 1000),
         pytest.param(0.1, 2),
         pytest.param(1000, 2),
-        pytest.param(0.1, 1000),
+        pytest.param(0.5, 1000),
     ],
 )
 def test_collate_requests(batch_timeout, batch_size):

tests/unit/test_lit_server.py

@@ -773,7 +773,8 @@ async def test_worker_restart_and_server_shutdown():
             resp = await ac.post("/predict", json={"input": 0}, timeout=2)
             assert resp.status_code == 500
 
-            await asyncio.sleep(2)
+            # Increased wait time for macOS where process spawning is slower
+            await asyncio.sleep(5 if sys.platform == "darwin" else 2)
 
             tasks = []
             for _ in range(50):

tests/unit/test_loops.py

@@ -173,13 +173,13 @@ def test_run_single_loop_with_async(async_loop_args, monkeypatch):
     with contextlib.suppress(KeyboardInterrupt):
         loop._run_single_loop_with_async(lit_api_mock, requests_queue, mock_transport, NOOP_CB_RUNNER)
 
-    response = asyncio.get_event_loop().run_until_complete(mock_transport.areceive(consumer_id=0))
+    response = asyncio.run(mock_transport.areceive(consumer_id=0))
     assert response == ("uuid-123", ((), "START", ls.utils.LoopResponseType.REGULAR, ANY))
-    response = asyncio.get_event_loop().run_until_complete(mock_transport.areceive(consumer_id=0))
+    response = asyncio.run(mock_transport.areceive(consumer_id=0))
     assert response == ("uuid-123", ({"output": 1}, ls.utils.LitAPIStatus.OK, ls.utils.LoopResponseType.REGULAR, ANY))
-    response = asyncio.get_event_loop().run_until_complete(mock_transport.areceive(consumer_id=1))
+    response = asyncio.run(mock_transport.areceive(consumer_id=1))
     assert response == ("uuid-234", ((), "START", ls.utils.LoopResponseType.REGULAR, ANY))
-    response = asyncio.get_event_loop().run_until_complete(mock_transport.areceive(consumer_id=1))
+    response = asyncio.run(mock_transport.areceive(consumer_id=1))
     assert response == ("uuid-234", ({"output": 4}, ls.utils.LitAPIStatus.OK, ls.utils.LoopResponseType.REGULAR, ANY))
 
 
@@ -292,7 +292,7 @@ def test_run_streaming_loop_with_async(mock_transport, monkeypatch):
         loop.run_streaming_loop_async(lit_api, requests_queue, mock_transport, NOOP_CB_RUNNER)
 
     for i in range(6):
-        response = asyncio.get_event_loop().run_until_complete(mock_transport.areceive(consumer_id=0))
+        response = asyncio.run(mock_transport.areceive(consumer_id=0))
         if i == 0:
             assert response == (
                 "uuid-123",