fix(functions): Move credential refresh to functions service init

Author: jonathanedey

firebase_admin/functions.py

@@ -24,7 +24,9 @@
 from dataclasses import dataclass
 
 from google.auth.compute_engine import Credentials as ComputeEngineCredentials
+from google.auth.exceptions import RefreshError
 from google.auth.transport import requests as google_auth_requests
+
 import requests
 import firebase_admin
 from firebase_admin import App
@@ -101,6 +103,12 @@ def __init__(self, app: App):
                 'GOOGLE_CLOUD_PROJECT environment variable.')
 
         self._credential = app.credential.get_credential()
+        try:
+            # Refresh the credential to ensure all attributes (e.g. service_account_email)
+            # are populated, preventing cold start errors.
+            self._credential.refresh(google_auth_requests.Request())
+        except RefreshError as err:
+            raise ValueError(f'Initial credential refresh failed: {err}') from err
         self._http_client = _http_client.JsonHttpClient(credential=self._credential)
 
     def task_queue(self, function_name: str, extension_id: Optional[str] = None) -> TaskQueue:
@@ -290,7 +298,6 @@ def _update_task_payload(self, task: Task, resource: Resource, extension_id: str
         # Meaning that it's credential should be a Compute Engine Credential.
         if _Validators.is_non_empty_string(extension_id) and \
             isinstance(self._credential, ComputeEngineCredentials):
-            self._credential.refresh(google_auth_requests.Request())
             id_token = self._credential.token
             task.http_request['headers'] = \
                 {**task.http_request['headers'], 'Authorization': f'Bearer {id_token}'}

tests/test_functions.py

@@ -17,7 +17,6 @@
 from datetime import datetime, timedelta, timezone
 import json
 import time
-from unittest import mock
 import pytest
 
 import firebase_admin
@@ -125,6 +124,10 @@ def test_task_enqueue(self):
         assert recorder[0].headers['x-goog-api-client'] == expected_metrics_header
         assert task_id == 'test-task-id'
 
+        task = json.loads(recorder[0].body.decode())['task']
+        assert task['http_request']['oidc_token'] == {'service_account_email': 'mock-email'}
+        assert task['http_request']['headers'] == {'Content-Type': 'application/json'}
+
     def test_task_enqueue_with_extension(self):
         resource_name = (
             'projects/test-project/locations/us-central1/queues/'
@@ -143,46 +146,68 @@ def test_task_enqueue_with_extension(self):
         assert recorder[0].headers['x-goog-api-client'] == expected_metrics_header
         assert task_id == 'test-task-id'
 
-    def test_task_delete(self):
-        _, recorder = self._instrument_functions_service()
-        queue = functions.task_queue('test-function-name')
-        queue.delete('test-task-id')
+        task = json.loads(recorder[0].body.decode())['task']
+        assert task['http_request']['oidc_token'] == {'service_account_email': 'mock-email'}
+        assert task['http_request']['headers'] == {'Content-Type': 'application/json'}
+
+    def test_task_enqueue_compute_engine(self):
+        app = firebase_admin.initialize_app(
+            testutils.MockComputeEngineCredential(),
+            options={'projectId': 'test-project'},
+            name='test-project-gce')
+        _, recorder = self._instrument_functions_service(app)
+        queue = functions.task_queue('test-function-name', app=app)
+        task_id = queue.enqueue(_DEFAULT_DATA)
         assert len(recorder) == 1
-        assert recorder[0].method == 'DELETE'
-        assert recorder[0].url == _DEFAULT_TASK_URL
-        expected_metrics_header = _utils.get_metrics_header() + ' mock-cred-metric-tag'
+        assert recorder[0].method == 'POST'
+        assert recorder[0].url == _DEFAULT_REQUEST_URL
+        assert recorder[0].headers['Content-Type'] == 'application/json'
+        assert recorder[0].headers['Authorization'] == 'Bearer mock-compute-engine-token'
+        expected_metrics_header = _utils.get_metrics_header() + ' mock-gce-cred-metric-tag'
         assert recorder[0].headers['x-goog-api-client'] == expected_metrics_header
+        assert task_id == 'test-task-id'
 
-    @mock.patch('firebase_admin.functions.isinstance')
-    def test_task_enqueue_with_extension_refreshes_credential(self, mock_isinstance):
-        # Force the code to take the ComputeEngineCredentials path
-        mock_isinstance.return_value = True
+        task = json.loads(recorder[0].body.decode())['task']
+        assert task['http_request']['oidc_token'] == {'service_account_email': 'mock-gce-email'}
+        assert task['http_request']['headers'] == {'Content-Type': 'application/json'}
 
-        # Create a custom response with the extension ID in the resource name
+    def test_task_enqueue_with_extension_compute_engine(self):
         resource_name = (
             'projects/test-project/locations/us-central1/queues/'
             'ext-test-extension-id-test-function-name/tasks'
         )
         extension_response = json.dumps({'name': resource_name + '/test-task-id'})
+        app = firebase_admin.initialize_app(
+            testutils.MockComputeEngineCredential(),
+            options={'projectId': 'test-project'},
+            name='test-project-gce-extensions')
+        _, recorder = self._instrument_functions_service(app, payload=extension_response)
+        queue = functions.task_queue('test-function-name', 'test-extension-id', app)
+        task_id = queue.enqueue(_DEFAULT_DATA)
+        assert len(recorder) == 1
+        assert recorder[0].method == 'POST'
+        assert recorder[0].url == _CLOUD_TASKS_URL + resource_name
+        assert recorder[0].headers['Content-Type'] == 'application/json'
+        assert recorder[0].headers['Authorization'] == 'Bearer mock-compute-engine-token'
+        expected_metrics_header = _utils.get_metrics_header() + ' mock-gce-cred-metric-tag'
+        assert recorder[0].headers['x-goog-api-client'] == expected_metrics_header
+        assert task_id == 'test-task-id'
 
-        # Instrument the service and get the underlying credential mock
-        functions_service, recorder = self._instrument_functions_service(payload=extension_response)
-        mock_credential = functions_service._credential
-        mock_credential.token = 'mock-id-token'
-        mock_credential.refresh = mock.MagicMock()
-
-        # Create a TaskQueue with an extension ID
-        queue = functions_service.task_queue('test-function-name', 'test-extension-id')
-
-        # Enqueue a task
-        queue.enqueue(_DEFAULT_DATA)
-
-        # Assert that the credential was refreshed
-        mock_credential.refresh.assert_called_once()
+        task = json.loads(recorder[0].body.decode())['task']
+        assert 'oidc_token' not in task['http_request']
+        assert task['http_request']['headers'] == {
+            'Content-Type': 'application/json',
+            'Authorization': 'Bearer mock-compute-engine-token'}
 
-        # Assert that the correct token was used in the header
+    def test_task_delete(self):
+        _, recorder = self._instrument_functions_service()
+        queue = functions.task_queue('test-function-name')
+        queue.delete('test-task-id')
         assert len(recorder) == 1
-        assert recorder[0].headers['Authorization'] == 'Bearer mock-id-token'
+        assert recorder[0].method == 'DELETE'
+        assert recorder[0].url == _DEFAULT_TASK_URL
+        expected_metrics_header = _utils.get_metrics_header() + ' mock-cred-metric-tag'
+        assert recorder[0].headers['x-goog-api-client'] == expected_metrics_header
 
 class TestTaskQueueOptions:
 

tests/testutils.py

@@ -118,10 +118,11 @@ class MockGoogleCredential(credentials.Credentials):
     """A mock Google authentication credential."""
     def refresh(self, request):
         self.token = 'mock-token'
+        self._service_account_email = "mock-email"
 
     @property
     def service_account_email(self):
-        return 'mock-email'
+        return self._service_account_email
 
     # Simulate x-goog-api-client modification in credential refresh
     def _metric_header_for_usage(self):
@@ -141,6 +142,10 @@ class MockGoogleComputeEngineCredential(compute_engine.Credentials):
     """A mock Compute Engine credential"""
     def refresh(self, request):
         self.token = 'mock-compute-engine-token'
+        self._service_account_email = 'mock-gce-email'
+
+    def _metric_header_for_usage(self):
+        return 'mock-gce-cred-metric-tag'
 
 class MockComputeEngineCredential(firebase_admin.credentials.Base):
     """A mock Firebase credential implementation."""