Skip streaming integration tests if REPLICATE_API_TOKEN isn't set

Signed-off-by: Mattt Zmuda <mattt@replicate.com>
replicate · Jun 28, 2024 · 16ac0d5 · 16ac0d5
1 parent b86aca7
commit 16ac0d5
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/tests/test_stream.py b/tests/test_stream.py
@@ -1,9 +1,16 @@
+import os
+
 import pytest
 
 import replicate
 from replicate.stream import ServerSentEvent
 
+skip_if_no_token = pytest.mark.skipif(
+    os.environ.get("REPLICATE_API_TOKEN") is None, reason="REPLICATE_API_TOKEN not set"
+)
+
 
+@skip_if_no_token
 @pytest.mark.asyncio
 @pytest.mark.parametrize("async_flag", [True, False])
 async def test_stream(async_flag, record_mode):
@@ -32,6 +39,7 @@ async def test_stream(async_flag, record_mode):
     assert any(event.event == ServerSentEvent.EventType.DONE for event in events)
 
 
+@skip_if_no_token
 @pytest.mark.asyncio
 @pytest.mark.parametrize("async_flag", [True, False])
 async def test_stream_prediction(async_flag, record_mode):