scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-03 02:22:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

test: update reset_key_spend test to match direct cache set

Browse source 
The implementation now sets spend_counter_cache directly instead of
calling _invalidate_spend_counter. Update the test to verify the
in_memory_cache.set_cache call with the correct key, value, and ttl.
This commit is contained in:
michaelxer 2026-06-11 02:32:28 +07:00
parent c56fdf2b66
commit 3d30cf0d3c
1 changed files with 92 additions and 84 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

176
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -6496,14 +6496,20 @@ async def test_reset_key_spend_success(monkeypatch):
patch(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object"
) as mock_delete_cache,
patch(
"litellm.proxy.proxy_server._invalidate_spend_counter"
) as mock_invalidate,
):
mock_hash_token.return_value = hashed_key
mock_check_admin.return_value = None
mock_delete_cache.return_value = None
# Mock spend_counter_cache to verify direct cache set instead of
# _invalidate_spend_counter (removed in favour of atomic cache write).
mock_spend_counter_cache = MagicMock()
mock_spend_counter_cache.redis_cache = None
monkeypatch.setattr(
"litellm.proxy.proxy_server.spend_counter_cache",
mock_spend_counter_cache,
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
@ -6523,7 +6529,9 @@ async def test_reset_key_spend_success(monkeypatch):
assert response["max_budget"] == 200.0
mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once()
mock_delete_cache.assert_awaited_once()
mock_invalidate.assert_awaited_once_with(counter_key=f"spend:key:{hashed_key}")
mock_spend_counter_cache.in_memory_cache.set_cache.assert_called_once_with(
key=f"spend:key:{hashed_key}", value=50.0, ttl=60
)
@pytest.mark.asyncio
@ -11742,83 +11750,83 @@ async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption(
assert str(code) == "400"
assert "cannot exceed" in msg.lower()
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_null_clears_fields():
"""
When budget_duration is explicitly set to null, prepare_key_update_data
should produce budget_duration=None and budget_reset_at=None so Prisma
clears them in the DB.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", budget_duration=None)
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert "budget_duration" in result
assert result["budget_duration"] is None
assert "budget_reset_at" in result
assert result["budget_reset_at"] is None
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_not_sent_excluded():
"""
When budget_duration is NOT sent in the request (unset), it should not
appear in the result dict at all — the existing DB value stays unchanged.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", models=["gpt-4"])
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert "budget_duration" not in result
assert "budget_reset_at" not in result
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_valid_sets_reset():
"""
When budget_duration is set to a valid duration string, both
budget_duration and budget_reset_at should be populated.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", budget_duration="30d")
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert result["budget_duration"] == "30d"
assert result["budget_reset_at"] is not None
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_null_clears_fields():
"""
When budget_duration is explicitly set to null, prepare_key_update_data
should produce budget_duration=None and budget_reset_at=None so Prisma
clears them in the DB.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", budget_duration=None)
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert "budget_duration" in result
assert result["budget_duration"] is None
assert "budget_reset_at" in result
assert result["budget_reset_at"] is None
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_not_sent_excluded():
"""
When budget_duration is NOT sent in the request (unset), it should not
appear in the result dict at all — the existing DB value stays unchanged.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", models=["gpt-4"])
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert "budget_duration" not in result
assert "budget_reset_at" not in result
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_valid_sets_reset():
"""
When budget_duration is set to a valid duration string, both
budget_duration and budget_reset_at should be populated.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", budget_duration="30d")
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert result["budget_duration"] == "30d"
assert result["budget_reset_at"] is not None