Skip to content

Commit 302f5da

Browse files
aronphilandstuff
authored andcommitted
Make scope lookup case insensitive
1 parent 52c373c commit 302f5da

File tree

2 files changed

+38
-24
lines changed

2 files changed

+38
-24
lines changed

replicate/client.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -355,10 +355,11 @@ def _get_api_token_from_environment() -> Optional[str]:
355355

356356
if hasattr(cog, "current_scope"):
357357
scope = cog.current_scope()
358-
if scope and hasattr(scope, "content") and isinstance(scope.content, dict):
359-
if "replicate_api_token" in scope.content:
360-
return scope.content["replicate_api_token"]
361-
except (ImportError, AttributeError, Exception):
358+
if scope and hasattr(scope, "context") and isinstance(scope.context, dict):
359+
for key, value in scope.context.items():
360+
if key.upper() == "REPLICATE_API_TOKEN":
361+
return scope.context[key]
362+
except: # noqa: S110,E722,BLE001 we don't want this code to cause clients to fail
362363
pass
363364

364365
return os.environ.get("REPLICATE_API_TOKEN")

tests/test_client.py

Lines changed: 33 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,7 @@ def test_cog_not_available_falls_back_to_env(self):
127127
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
128128
with mock.patch.dict(sys.modules, {"cog": None}):
129129
token = _get_api_token_from_environment()
130-
assert token == "env-token"
130+
assert token == "env-token" # noqa: S105
131131

132132
def test_cog_import_error_falls_back_to_env(self):
133133
"""Test fallback to environment when cog import raises exception."""
@@ -137,7 +137,7 @@ def test_cog_import_error_falls_back_to_env(self):
137137
side_effect=ModuleNotFoundError("No module named 'cog'"),
138138
):
139139
token = _get_api_token_from_environment()
140-
assert token == "env-token"
140+
assert token == "env-token" # noqa: S105
141141

142142
def test_cog_no_current_scope_method_falls_back_to_env(self):
143143
"""Test fallback when cog exists but has no current_scope method."""
@@ -147,7 +147,7 @@ def test_cog_no_current_scope_method_falls_back_to_env(self):
147147
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
148148
with mock.patch.dict(sys.modules, {"cog": mock_cog}):
149149
token = _get_api_token_from_environment()
150-
assert token == "env-token"
150+
assert token == "env-token" # noqa: S105
151151

152152
def test_cog_current_scope_returns_none_falls_back_to_env(self):
153153
"""Test fallback when current_scope() returns None."""
@@ -157,64 +157,77 @@ def test_cog_current_scope_returns_none_falls_back_to_env(self):
157157
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
158158
with mock.patch.dict(sys.modules, {"cog": mock_cog}):
159159
token = _get_api_token_from_environment()
160-
assert token == "env-token"
160+
assert token == "env-token" # noqa: S105
161161

162-
def test_cog_scope_no_content_attr_falls_back_to_env(self):
163-
"""Test fallback when scope has no content attribute."""
162+
def test_cog_scope_no_context_attr_falls_back_to_env(self):
163+
"""Test fallback when scope has no context attribute."""
164164
mock_scope = mock.MagicMock()
165-
del mock_scope.content # Remove the content attribute
165+
del mock_scope.context # Remove the context attribute
166166

167167
mock_cog = mock.MagicMock()
168168
mock_cog.current_scope.return_value = mock_scope
169169

170170
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
171171
with mock.patch.dict(sys.modules, {"cog": mock_cog}):
172172
token = _get_api_token_from_environment()
173-
assert token == "env-token"
173+
assert token == "env-token" # noqa: S105
174174

175-
def test_cog_scope_content_not_dict_falls_back_to_env(self):
176-
"""Test fallback when scope.content is not a dictionary."""
175+
def test_cog_scope_context_not_dict_falls_back_to_env(self):
176+
"""Test fallback when scope.context is not a dictionary."""
177177
mock_scope = mock.MagicMock()
178-
mock_scope.content = "not a dict"
178+
mock_scope.context = "not a dict"
179179

180180
mock_cog = mock.MagicMock()
181181
mock_cog.current_scope.return_value = mock_scope
182182

183183
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
184184
with mock.patch.dict(sys.modules, {"cog": mock_cog}):
185185
token = _get_api_token_from_environment()
186-
assert token == "env-token"
186+
assert token == "env-token" # noqa: S105
187187

188188
def test_cog_scope_no_replicate_api_token_key_falls_back_to_env(self):
189-
"""Test fallback when replicate_api_token key is missing from content."""
189+
"""Test fallback when replicate_api_token key is missing from context."""
190190
mock_scope = mock.MagicMock()
191-
mock_scope.content = {"other_key": "other_value"} # Missing replicate_api_token
191+
mock_scope.context = {"other_key": "other_value"} # Missing replicate_api_token
192192

193193
mock_cog = mock.MagicMock()
194194
mock_cog.current_scope.return_value = mock_scope
195195

196196
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
197197
with mock.patch.dict(sys.modules, {"cog": mock_cog}):
198198
token = _get_api_token_from_environment()
199-
assert token == "env-token"
199+
assert token == "env-token" # noqa: S105
200200

201201
def test_cog_scope_replicate_api_token_valid_string(self):
202202
"""Test successful retrieval of non-empty token from cog."""
203203
mock_scope = mock.MagicMock()
204-
mock_scope.content = {"replicate_api_token": "cog-token"}
204+
mock_scope.context = {"REPLICATE_API_TOKEN": "cog-token"}
205205

206206
mock_cog = mock.MagicMock()
207207
mock_cog.current_scope.return_value = mock_scope
208208

209209
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
210210
with mock.patch.dict(sys.modules, {"cog": mock_cog}):
211211
token = _get_api_token_from_environment()
212-
assert token == "cog-token"
212+
assert token == "cog-token" # noqa: S105
213+
214+
def test_cog_scope_replicate_api_token_case_insensitive(self):
215+
"""Test successful retrieval of non-empty token from cog ignoring case."""
216+
mock_scope = mock.MagicMock()
217+
mock_scope.context = {"replicate_api_token": "cog-token"}
218+
219+
mock_cog = mock.MagicMock()
220+
mock_cog.current_scope.return_value = mock_scope
221+
222+
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
223+
with mock.patch.dict(sys.modules, {"cog": mock_cog}):
224+
token = _get_api_token_from_environment()
225+
assert token == "cog-token" # noqa: S105
213226

214227
def test_cog_scope_replicate_api_token_empty_string(self):
215228
"""Test that empty string from cog is returned (not falling back to env)."""
216229
mock_scope = mock.MagicMock()
217-
mock_scope.content = {"replicate_api_token": ""} # Empty string
230+
mock_scope.context = {"replicate_api_token": ""} # Empty string
218231

219232
mock_cog = mock.MagicMock()
220233
mock_cog.current_scope.return_value = mock_scope
@@ -227,7 +240,7 @@ def test_cog_scope_replicate_api_token_empty_string(self):
227240
def test_cog_scope_replicate_api_token_none(self):
228241
"""Test that None from cog is returned (not falling back to env)."""
229242
mock_scope = mock.MagicMock()
230-
mock_scope.content = {"replicate_api_token": None}
243+
mock_scope.context = {"replicate_api_token": None}
231244

232245
mock_cog = mock.MagicMock()
233246
mock_cog.current_scope.return_value = mock_scope
@@ -245,7 +258,7 @@ def test_cog_current_scope_raises_exception_falls_back_to_env(self):
245258
with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "env-token"}):
246259
with mock.patch.dict(sys.modules, {"cog": mock_cog}):
247260
token = _get_api_token_from_environment()
248-
assert token == "env-token"
261+
assert token == "env-token" # noqa: S105
249262

250263
def test_no_env_token_returns_none(self):
251264
"""Test that None is returned when no environment token is set and cog unavailable."""

0 commit comments

Comments
 (0)