chore(weave): Add pytest retry (#2919)

wandb · Nov 6, 2024 · 9407a60 · 9407a60
1 parent 5ce195a
commit 9407a60
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 3 deletions.
diff --git a/noxfile.py b/noxfile.py
@@ -98,6 +98,7 @@ def tests(session, shard):
 
     session.run(
         "pytest",
+        "--strict-markers",
         "--cov=weave",
         "--cov-report=html",
         "--cov-branch",

diff --git a/pyproject.toml b/pyproject.toml
@@ -86,6 +86,7 @@ test = [
   "pytest-cov>=5.0.0",
   "pytest-xdist>=3.1.0",
   "pytest-rerunfailures>=12.0",
+  "pytest-rerunfailures>=14.0",
   "clickhouse_connect==0.7.0",
   "fastapi>=0.110.0",
   "sqlparse==0.5.0",

diff --git a/tests/integrations/google_ai_studio/google_ai_studio_test.py b/tests/integrations/google_ai_studio/google_ai_studio_test.py
@@ -47,7 +47,7 @@ def assert_correct_summary(summary: dict, trace_name: str):
     assert summary["weave"]["latency_ms"] > 0
 
 
-@pytest.mark.retry(max_attempts=5, delay=5)
+@pytest.mark.flaky(reruns=5, reruns_delay=2)
 @pytest.mark.skip_clickhouse_client
 def test_content_generation(client):
     import google.generativeai as genai
@@ -69,7 +69,7 @@ def test_content_generation(client):
     assert_correct_summary(call.summary, trace_name)
 
 
-@pytest.mark.retry(max_attempts=5, delay=5)
+@pytest.mark.flaky(reruns=5, reruns_delay=2)
 @pytest.mark.skip_clickhouse_client
 def test_content_generation_stream(client):
     import google.generativeai as genai
@@ -95,7 +95,7 @@ def test_content_generation_stream(client):
     assert_correct_summary(call.summary, trace_name)
 
 
-@pytest.mark.retry(max_attempts=5, delay=5)
+@pytest.mark.flaky(reruns=5, reruns_delay=2)
 @pytest.mark.asyncio
 @pytest.mark.skip_clickhouse_client
 async def test_content_generation_async(client):