Skip to content

Commit 2df34ed

Browse files
committed
Silence warning when using cog.current_scope()
1 parent e8acdb2 commit 2df34ed

File tree

2 files changed

+18
-0
lines changed

2 files changed

+18
-0
lines changed

replicate/client.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -352,6 +352,11 @@ def _get_api_token_from_environment() -> Optional[str]:
352352
"""Get API token from cog current scope if available, otherwise from environment."""
353353
try:
354354
import cog # noqa: I001 # pyright: ignore [reportMissingImports]
355+
import warnings
356+
357+
warnings.filterwarnings(
358+
"ignore", message="current_scope", category=cog.ExperimentalFeatureWarning
359+
)
355360

356361
for key, value in cog.current_scope().context.items():
357362
if key.upper() == "REPLICATE_API_TOKEN":

tests/test_client.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,9 @@ def mock_send(request):
119119
mock_send_wrapper.assert_called_once()
120120

121121

122+
class ExperimentalFeatureWarning(Warning): ...
123+
124+
122125
class TestGetApiToken:
123126
"""Test cases for _get_api_token_from_environment function covering all import paths."""
124127

@@ -142,6 +145,7 @@ def test_cog_import_error_falls_back_to_env(self):
142145
def test_cog_no_current_scope_method_falls_back_to_env(self):
143146
"""Test fallback when cog exists but has no current_scope method."""
144147
mock_cog = mock.MagicMock()
148+
mock_cog.ExperimentalFeatureWarning = ExperimentalFeatureWarning
145149
del mock_cog.current_scope # Remove the method
146150

147151
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
@@ -152,6 +156,7 @@ def test_cog_no_current_scope_method_falls_back_to_env(self):
152156
def test_cog_current_scope_returns_none_falls_back_to_env(self):
153157
"""Test fallback when current_scope() returns None."""
154158
mock_cog = mock.MagicMock()
159+
mock_cog.ExperimentalFeatureWarning = ExperimentalFeatureWarning
155160
mock_cog.current_scope.return_value = None
156161

157162
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
@@ -165,6 +170,7 @@ def test_cog_scope_no_context_attr_falls_back_to_env(self):
165170
del mock_scope.context # Remove the context attribute
166171

167172
mock_cog = mock.MagicMock()
173+
mock_cog.ExperimentalFeatureWarning = ExperimentalFeatureWarning
168174
mock_cog.current_scope.return_value = mock_scope
169175

170176
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
@@ -178,6 +184,7 @@ def test_cog_scope_context_not_dict_falls_back_to_env(self):
178184
mock_scope.context = "not a dict"
179185

180186
mock_cog = mock.MagicMock()
187+
mock_cog.ExperimentalFeatureWarning = ExperimentalFeatureWarning
181188
mock_cog.current_scope.return_value = mock_scope
182189

183190
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
@@ -191,6 +198,7 @@ def test_cog_scope_no_replicate_api_token_key_falls_back_to_env(self):
191198
mock_scope.context = {"other_key": "other_value"} # Missing replicate_api_token
192199

193200
mock_cog = mock.MagicMock()
201+
mock_cog.ExperimentalFeatureWarning = ExperimentalFeatureWarning
194202
mock_cog.current_scope.return_value = mock_scope
195203

196204
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
@@ -204,6 +212,7 @@ def test_cog_scope_replicate_api_token_valid_string(self):
204212
mock_scope.context = {"REPLICATE_API_TOKEN": "cog-token"}
205213

206214
mock_cog = mock.MagicMock()
215+
mock_cog.ExperimentalFeatureWarning = ExperimentalFeatureWarning
207216
mock_cog.current_scope.return_value = mock_scope
208217

209218
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
@@ -217,6 +226,7 @@ def test_cog_scope_replicate_api_token_case_insensitive(self):
217226
mock_scope.context = {"replicate_api_token": "cog-token"}
218227

219228
mock_cog = mock.MagicMock()
229+
mock_cog.ExperimentalFeatureWarning = ExperimentalFeatureWarning
220230
mock_cog.current_scope.return_value = mock_scope
221231

222232
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
@@ -230,6 +240,7 @@ def test_cog_scope_replicate_api_token_empty_string(self):
230240
mock_scope.context = {"replicate_api_token": ""} # Empty string
231241

232242
mock_cog = mock.MagicMock()
243+
mock_cog.ExperimentalFeatureWarning = ExperimentalFeatureWarning
233244
mock_cog.current_scope.return_value = mock_scope
234245

235246
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
@@ -243,6 +254,7 @@ def test_cog_scope_replicate_api_token_none(self):
243254
mock_scope.context = {"replicate_api_token": None}
244255

245256
mock_cog = mock.MagicMock()
257+
mock_cog.ExperimentalFeatureWarning = ExperimentalFeatureWarning
246258
mock_cog.current_scope.return_value = mock_scope
247259

248260
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
@@ -253,6 +265,7 @@ def test_cog_scope_replicate_api_token_none(self):
253265
def test_cog_current_scope_raises_exception_falls_back_to_env(self):
254266
"""Test fallback when current_scope() raises an exception."""
255267
mock_cog = mock.MagicMock()
268+
mock_cog.ExperimentalFeatureWarning = ExperimentalFeatureWarning
256269
mock_cog.current_scope.side_effect = RuntimeError("Scope error")
257270

258271
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):

0 commit comments

Comments
 (0)