Inject mock API token into tests

Signed-off-by: Mattt Zmuda <[email protected]>

Author: mattt

tests/conftest.py

@@ -1,8 +1,18 @@
 import os
+from unittest import mock
 
 import pytest
 
 
+@pytest.fixture(scope="session")
+def mock_replicate_api_token(scope="class"):
+    if os.environ.get("REPLICATE_API_TOKEN", "") != "":
+        yield
+    else:
+        with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "test-token"}):
+            yield
+
+
 @pytest.fixture(scope="module")
 def vcr_config():
     return {"allowed_hosts": ["api.replicate.com"], "filter_headers": ["authorization"]}

tests/test_model.py

@@ -5,7 +5,7 @@
 
 @pytest.mark.vcr("models-get.yaml")
 @pytest.mark.asyncio
-async def test_models_get():
+async def test_models_get(mock_replicate_api_token):
     model = replicate.models.get("stability-ai/sdxl")
 
     assert model is not None

tests/test_prediction.py

@@ -5,7 +5,7 @@
 
 @pytest.mark.vcr("predictions-create.yaml")
 @pytest.mark.asyncio
-async def test_predictions_create():
+async def test_predictions_create(mock_replicate_api_token):
     input = {
         "prompt": "a studio photo of a rainbow colored corgi",
         "width": 512,
@@ -29,7 +29,7 @@ async def test_predictions_create():
 
 @pytest.mark.vcr("predictions-get.yaml")
 @pytest.mark.asyncio
-async def test_predictions_get():
+async def test_predictions_get(mock_replicate_api_token):
     id = "vgcm4plb7tgzlyznry5d5jkgvu"
 
     prediction = replicate.predictions.get(id)
@@ -39,7 +39,7 @@ async def test_predictions_get():
 
 @pytest.mark.vcr("predictions-cancel.yaml")
 @pytest.mark.asyncio
-async def test_predictions_cancel():
+async def test_predictions_cancel(mock_replicate_api_token):
     input = {
         "prompt": "a studio photo of a rainbow colored corgi",
         "width": 512,

tests/test_run.py

@@ -6,7 +6,7 @@
 
 @pytest.mark.vcr("run.yaml")
 @pytest.mark.asyncio
-async def test_run():
+async def test_run(mock_replicate_api_token):
     version = "a00d0b7dcbb9c3fbb34ba87d2d5b46c56969c84a628bf778a7fdaec30b1b99c5"
 
     input = {
@@ -28,6 +28,6 @@ async def test_run():
 
 
 @pytest.mark.vcr
-def test_run_with_invalid_identifier():
+def test_run_with_invalid_identifier(mock_replicate_api_token):
     with pytest.raises(ReplicateError):
         replicate.run("invalid")

tests/test_training.py

@@ -8,7 +8,7 @@
 
 @pytest.mark.vcr("trainings-create.yaml")
 @pytest.mark.asyncio
-async def test_trainings_create():
+async def test_trainings_create(mock_replicate_api_token):
     training = replicate.trainings.create(
         "stability-ai/sdxl:a00d0b7dcbb9c3fbb34ba87d2d5b46c56969c84a628bf778a7fdaec30b1b99c5",
         input={
@@ -24,7 +24,7 @@ async def test_trainings_create():
 
 @pytest.mark.vcr("trainings-create__invalid-destination.yaml")
 @pytest.mark.asyncio
-async def test_trainings_create_with_invalid_destination():
+async def test_trainings_create_with_invalid_destination(mock_replicate_api_token):
     with pytest.raises(ReplicateException):
         replicate.trainings.create(
             "stability-ai/sdxl:a00d0b7dcbb9c3fbb34ba87d2d5b46c56969c84a628bf778a7fdaec30b1b99c5",
@@ -37,7 +37,7 @@ async def test_trainings_create_with_invalid_destination():
 
 @pytest.mark.vcr("trainings-get.yaml")
 @pytest.mark.asyncio
-async def test_trainings_get():
+async def test_trainings_get(mock_replicate_api_token):
     id = "ckcbvmtbvg6di3b3uhvccytnfm"
 
     training = replicate.trainings.get(id)
@@ -48,7 +48,7 @@ async def test_trainings_get():
 
 @pytest.mark.vcr("trainings-cancel.yaml")
 @pytest.mark.asyncio
-async def test_trainings_cancel():
+async def test_trainings_cancel(mock_replicate_api_token):
     input = {
         "input_images": input_images_url,
         "use_face_detection_instead": True,