diff --git a/litellm/integrations/otel/runtime.py b/litellm/integrations/otel/runtime.py index 0455faf8ce6..5e10cae079f 100644 --- a/litellm/integrations/otel/runtime.py +++ b/litellm/integrations/otel/runtime.py @@ -100,9 +100,7 @@ def traced( def _decorator(func: Callable) -> Callable: if not inspect.iscoroutinefunction(func): - raise TypeError( - f"@traced requires an async function; got {func!r}" - ) + raise TypeError(f"@traced requires an async function; got {func!r}") @functools.wraps(func) async def _wrapper(*args, **kwargs): diff --git a/tests/test_litellm/identity/test_invalidation.py b/tests/test_litellm/identity/test_invalidation.py index e21e00cf34e..c36efc5cb2b 100644 --- a/tests/test_litellm/identity/test_invalidation.py +++ b/tests/test_litellm/identity/test_invalidation.py @@ -24,9 +24,7 @@ async def test_token_invalidation_drops_entry(): await cache.set(uak.token, uak) assert await cache.get(uak.token) is not None - await invalidate_identity_for_token( - token_hash=uak.token, dual_cache=backend - ) + await invalidate_identity_for_token(token_hash=uak.token, dual_cache=backend) assert await cache.get(uak.token) is None @@ -37,9 +35,7 @@ async def test_team_invalidation_bumps_generation(): uak = UserAPIKeyAuth(api_key="sk-x", user_id="u1", team_id="t-rotate") await cache.set(uak.token, uak) - await invalidate_identity_for_team( - team_id="t-rotate", dual_cache=backend - ) + await invalidate_identity_for_team(team_id="t-rotate", dual_cache=backend) assert await cache.get(uak.token) is None @@ -52,9 +48,7 @@ async def test_user_invalidation_is_scoped_to_user(): await cache.set(uak_a.token, uak_a) await cache.set(uak_b.token, uak_b) - await invalidate_identity_for_user( - user_id="u-rotate", dual_cache=backend - ) + await invalidate_identity_for_user(user_id="u-rotate", dual_cache=backend) assert await cache.get(uak_a.token) is None assert await cache.get(uak_b.token) is not None @@ -67,7 +61,5 @@ async def test_org_invalidation_drops_org_scoped_entry(): uak = UserAPIKeyAuth(api_key="sk-x", user_id="u1", org_id="org-rotate") await cache.set(uak.token, uak) - await invalidate_identity_for_org( - org_id="org-rotate", dual_cache=backend - ) + await invalidate_identity_for_org(org_id="org-rotate", dual_cache=backend) assert await cache.get(uak.token) is None diff --git a/tests/test_litellm/identity/test_observability.py b/tests/test_litellm/identity/test_observability.py index c576061efeb..fc873115973 100644 --- a/tests/test_litellm/identity/test_observability.py +++ b/tests/test_litellm/identity/test_observability.py @@ -69,6 +69,7 @@ async def test_db_call_role_does_not_open_a_new_span(): def test_decorator_rejects_sync_functions(): with pytest.raises(TypeError): + @traced("identity.sync-not-allowed", role=SpanRole.SERVICE) def sync_fn(): return None diff --git a/tests/test_litellm/identity/test_resolver.py b/tests/test_litellm/identity/test_resolver.py index 63b05ae1ba3..4f9cf8b693f 100644 --- a/tests/test_litellm/identity/test_resolver.py +++ b/tests/test_litellm/identity/test_resolver.py @@ -27,9 +27,7 @@ def _fake_request(headers=None, client_host=None): def _jwt(claims): def b(d): - return ( - base64.urlsafe_b64encode(json.dumps(d).encode()).rstrip(b"=").decode() - ) + return base64.urlsafe_b64encode(json.dumps(d).encode()).rstrip(b"=").decode() return f"{b({'alg':'HS256','typ':'JWT'})}.{b(claims)}.sig" diff --git a/tests/test_litellm/identity/test_store.py b/tests/test_litellm/identity/test_store.py index 6f9135b9ceb..b32f35f91ad 100644 --- a/tests/test_litellm/identity/test_store.py +++ b/tests/test_litellm/identity/test_store.py @@ -45,9 +45,7 @@ async def test_cache_hit_skips_db(): prisma = _stub_prisma_client() from litellm.proxy._types import UserAPIKeyAuth - seed = UserAPIKeyAuth( - token="hash-cached", user_id="u-cache", team_id="t-cache" - ) + seed = UserAPIKeyAuth(token="hash-cached", user_id="u-cache", team_id="t-cache") await identity_cache.set("hash-cached", seed) result = await load_identity( @@ -71,9 +69,7 @@ async def test_cache_miss_hits_db_once_and_caches(): token="hash-db", user_id="u-db", team_id="t-db" ) - with patch( - "litellm.identity.store._populate_legacy_cache", new=AsyncMock() - ): + with patch("litellm.identity.store._populate_legacy_cache", new=AsyncMock()): result = await load_identity( hashed_token="hash-db", prisma_client=prisma, @@ -97,9 +93,7 @@ async def test_warm_load_after_cold_does_not_hit_db(): token="hash-warm", user_id="u-warm", team_id="t-warm" ) - with patch( - "litellm.identity.store._populate_legacy_cache", new=AsyncMock() - ): + with patch("litellm.identity.store._populate_legacy_cache", new=AsyncMock()): await load_identity( hashed_token="hash-warm", prisma_client=prisma, @@ -146,9 +140,7 @@ async def test_bundled_user_survives_cache_roundtrip_as_typed_model(): uak = UserAPIKeyAuth( api_key="sk-bundle", user_id="u-bundle", - user=LiteLLM_UserTable( - user_id="u-bundle", user_email="x@y.com", tpm_limit=10 - ), + user=LiteLLM_UserTable(user_id="u-bundle", user_email="x@y.com", tpm_limit=10), ) await identity_cache.set(uak.token, uak) got = await identity_cache.get(uak.token) @@ -177,11 +169,12 @@ async def test_cold_path_bundles_user_into_cache(): user_id=user_id, user_email="bundle@litellm.io", tpm_limit=42 ) - with patch( - "litellm.identity.store._populate_legacy_cache", new=AsyncMock() - ), patch( - "litellm.proxy.auth.auth_checks.get_user_object", - side_effect=_fake_get_user_object, + with ( + patch("litellm.identity.store._populate_legacy_cache", new=AsyncMock()), + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + side_effect=_fake_get_user_object, + ), ): result = await load_identity( hashed_token="hash-user", @@ -206,9 +199,7 @@ async def test_hydrated_copy_is_request_scoped(): token="hash-copy", user_id="u-copy", team_id="t-copy" ) - with patch( - "litellm.identity.store._populate_legacy_cache", new=AsyncMock() - ): + with patch("litellm.identity.store._populate_legacy_cache", new=AsyncMock()): a = await load_identity( hashed_token="hash-copy", prisma_client=prisma,