From 7d18bc289a1fa271f91d6f6cb29a1b6acaae8408 Mon Sep 17 00:00:00 2001 From: Louis Vauterin Date: Mon, 7 Sep 2026 13:09:14 +0200 Subject: [PATCH] test(jwt-key-mapping): restore the cache eviction tests dropped in the merge The merge with litellm_internal_staging resolved the test file by taking this branch's copy whole, which discarded the two tests staging had appended to the same end-of-file region: test_delete_evicts_cache_after_row_is_gone test_update_evicts_old_and_new_cache_keys_after_write Both sides only appended, so the conflict was additive and nothing had to be chosen between them. The eviction code those tests cover did survive in jwt_key_mapping_endpoints.py, so the branch was shipping it untested. Appending staging's block restores them alongside the nine resolver tests here. The file is now a superset of both sides, verified line by line, and the test quality count stays at the baseline of 28 because staging's block carries its own test-quality-ok suppressions. --- .../proxy_unit_tests/test_jwt_key_mapping.py | 83 +++++++++++++++++++ 1 file changed, 83 insertions(+) diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index 7620f368f25..10916325728 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -1724,3 +1724,86 @@ def test_update_rejects_a_token_that_is_not_a_sha256_hash(): assert exc_info.value.status_code == 400 assert "SHA-256" in exc_info.value.detail + + +# ────────────────────────────────────────────── +# Tests: cache eviction must happen AFTER the DB write commits +# ────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_delete_evicts_cache_after_row_is_gone(): + """A JWT request racing the delete must not keep the removed mapping authorized. + + The DB delete simulates a concurrent request re-caching the mapping mid-write. + If the endpoint evicts before the delete commits, that repopulated entry + survives until TTL and the deleted mapping stays usable. + """ + from litellm.proxy._types import DeleteJWTKeyMappingRequest + from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key + + cache_key = jwt_key_mapping_cache_key("email", "user@example.com") + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache(key=cache_key, value="hashed_token") + + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = _mock_mapping() + + async def concurrent_reader_repopulates(**kwargs): + await user_api_key_cache.async_set_cache(key=cache_key, value="hashed_token") + return _mock_mapping() + + mock_prisma.db.litellm_jwtkeymapping.delete.side_effect = concurrent_reader_repopulates + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), # test-quality-ok: proxy_server module global is the endpoint's only injection point + ): + result = await delete_jwt_key_mapping( + data=DeleteJWTKeyMappingRequest(id="mapping-1"), + user_api_key_dict=_make_admin_auth(), + ) + + assert result == {"status": "success"} + assert await user_api_key_cache.async_get_cache(cache_key) is None + + +@pytest.mark.asyncio +async def test_update_evicts_old_and_new_cache_keys_after_write(): + """Renaming a mapping's claim must leave neither claim serving stale cache. + + The DB update simulates a concurrent request re-caching the OLD mapping + mid-write. Both the old claim's entry (would restore the pre-rename token) + and the new claim's __NO_MAPPING__ sentinel (would 403 the renamed claim) + must be gone once the endpoint returns. + """ + from litellm.proxy._types import UpdateJWTKeyMappingRequest + from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key + + old_cache_key = jwt_key_mapping_cache_key("email", "user@example.com") + new_cache_key = jwt_key_mapping_cache_key("email", "renamed@example.com") + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache(key=old_cache_key, value="hashed_token") + await user_api_key_cache.async_set_cache(key=new_cache_key, value="__NO_MAPPING__") + + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = _mock_mapping() + + async def concurrent_reader_repopulates(**kwargs): + await user_api_key_cache.async_set_cache(key=old_cache_key, value="hashed_token") + return _mock_mapping(claim_value="renamed@example.com") + + mock_prisma.db.litellm_jwtkeymapping.update.side_effect = concurrent_reader_repopulates + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), # test-quality-ok: proxy_server module global is the endpoint's only injection point + ): + result = await update_jwt_key_mapping( + data=UpdateJWTKeyMappingRequest(id="mapping-1", jwt_claim_value="renamed@example.com"), + user_api_key_dict=_make_admin_auth(), + ) + + assert result.jwt_claim_value == "renamed@example.com" + assert await user_api_key_cache.async_get_cache(old_cache_key) is None + assert await user_api_key_cache.async_get_cache(new_cache_key) is None