Internal

checkpoint/orbax/checkpoint/experimental/tiering_service/auth.py

@@ -19,6 +19,7 @@
 """
 
 from collections.abc import Collection
+
 from absl import logging
 import grpc
 from orbax.checkpoint.experimental.tiering_service import db_schema
@@ -39,7 +40,8 @@ async def get_oauth_token(context: grpc.aio.ServicerContext) -> str | None:
     The extracted OAuth token string, or None if not found or malformed.
   """
   logging.debug("Extracting OAuth token from metadata")
-  metadata = dict(await context.invocation_metadata())  # pyrefly: ignore[not-async]
+  raw_metadata = context.invocation_metadata()
+  metadata = dict(raw_metadata)
   # Standard header for OAuth tokens in gRPC is 'authorization'.
   auth_header = metadata.get("authorization")
 

checkpoint/orbax/checkpoint/experimental/tiering_service/auth_test.py

@@ -28,7 +28,7 @@ async def test_get_oauth_token_success(self):
     context = mock.create_autospec(
         grpc.aio.ServicerContext, instance=True, spec_set=True
     )
-    context.invocation_metadata = mock.AsyncMock(
+    context.invocation_metadata = mock.Mock(
         return_value=(("authorization", "Bearer valid-token"),)
     )
     token = await auth.get_oauth_token(context)
@@ -38,15 +38,15 @@ async def test_get_oauth_token_no_header(self):
     context = mock.create_autospec(
         grpc.aio.ServicerContext, instance=True, spec_set=True
     )
-    context.invocation_metadata = mock.AsyncMock(return_value=())
+    context.invocation_metadata = mock.Mock(return_value=())
     token = await auth.get_oauth_token(context)
     self.assertIsNone(token)
 
   async def test_get_oauth_token_malformed_header(self):
     context = mock.create_autospec(
         grpc.aio.ServicerContext, instance=True, spec_set=True
     )
-    context.invocation_metadata = mock.AsyncMock(
+    context.invocation_metadata = mock.Mock(
         return_value=(("authorization", "not-bearer-token"),)
     )
     token = await auth.get_oauth_token(context)