Improve CI: retry flaky tests (#379)

* retry tests * update * update
Lightning-AI · Dec 3, 2024 · 103bd70 · 103bd70
1 parent 43d9928
commit 103bd70
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/_requirements/test.txt b/_requirements/test.txt
@@ -14,3 +14,4 @@ transformers
 openai>=1.12.0
 pillow
 numpy <2.0
+pytest-retry>=1.6.3
diff --git a/tests/test_simple.py b/tests/test_simple.py
@@ -143,6 +143,7 @@ def unbatch(self, output):
         return list(output)
 
 
+@pytest.mark.flaky(retries=3)
 @pytest.mark.asyncio
 async def test_timeout():
     # Scenario: first request completes, second request times out in queue
@@ -172,6 +173,7 @@ async def test_timeout():
             )
 
 
+@pytest.mark.flaky(retries=3)
 @pytest.mark.asyncio
 async def test_batch_timeout():
     # Scenario: first 2 requests finish as a batch and third request times out in queue