|
17 | 17 | from datetime import datetime, timedelta, timezone
|
18 | 18 | import json
|
19 | 19 | import time
|
| 20 | +from unittest import mock |
20 | 21 | import pytest
|
21 | 22 |
|
22 | 23 | import firebase_admin
|
@@ -152,6 +153,37 @@ def test_task_delete(self):
|
152 | 153 | expected_metrics_header = _utils.get_metrics_header() + ' mock-cred-metric-tag'
|
153 | 154 | assert recorder[0].headers['x-goog-api-client'] == expected_metrics_header
|
154 | 155 |
|
| 156 | + @mock.patch('firebase_admin.functions.isinstance') |
| 157 | + def test_task_enqueue_with_extension_refreshes_credential(self, mock_isinstance): |
| 158 | + # Force the code to take the ComputeEngineCredentials path |
| 159 | + mock_isinstance.return_value = True |
| 160 | + |
| 161 | + # Create a custom response with the extension ID in the resource name |
| 162 | + resource_name = ( |
| 163 | + 'projects/test-project/locations/us-central1/queues/' |
| 164 | + 'ext-test-extension-id-test-function-name/tasks' |
| 165 | + ) |
| 166 | + extension_response = json.dumps({'name': resource_name + '/test-task-id'}) |
| 167 | + |
| 168 | + # Instrument the service and get the underlying credential mock |
| 169 | + functions_service, recorder = self._instrument_functions_service(payload=extension_response) |
| 170 | + mock_credential = functions_service._credential |
| 171 | + mock_credential.token = 'mock-id-token' |
| 172 | + mock_credential.refresh = mock.MagicMock() |
| 173 | + |
| 174 | + # Create a TaskQueue with an extension ID |
| 175 | + queue = functions_service.task_queue('test-function-name', 'test-extension-id') |
| 176 | + |
| 177 | + # Enqueue a task |
| 178 | + queue.enqueue(_DEFAULT_DATA) |
| 179 | + |
| 180 | + # Assert that the credential was refreshed |
| 181 | + mock_credential.refresh.assert_called_once() |
| 182 | + |
| 183 | + # Assert that the correct token was used in the header |
| 184 | + assert len(recorder) == 1 |
| 185 | + assert recorder[0].headers['Authorization'] == 'Bearer mock-id-token' |
| 186 | + |
155 | 187 | class TestTaskQueueOptions:
|
156 | 188 |
|
157 | 189 | _DEFAULT_TASK_OPTS = {'schedule_delay_seconds': None, 'schedule_time': None, \
|
|
0 commit comments