|
7 | 7 | from importlib.metadata import version |
8 | 8 | import threading |
9 | 9 | import time |
| 10 | +from typing import Any |
10 | 11 |
|
11 | 12 | import grpc |
12 | 13 | from azure.core.credentials import AccessToken |
@@ -39,6 +40,21 @@ def GetInstance(self, request, context): |
39 | 40 | return response |
40 | 41 |
|
41 | 42 |
|
| 43 | +class _TestTokenCredential: |
| 44 | + """Minimal TokenCredential stub that counts how often a token is requested.""" |
| 45 | + |
| 46 | + def __init__(self): |
| 47 | + self._lock = threading.Lock() |
| 48 | + self.calls = 0 |
| 49 | + |
| 50 | + def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken: |
| 51 | + with self._lock: |
| 52 | + self.calls += 1 |
| 53 | + call_number = self.calls |
| 54 | + time.sleep(0.02) |
| 55 | + return AccessToken(f"token-{call_number}", int(time.time()) + 3600) |
| 56 | + |
| 57 | + |
42 | 58 | class TestDurableTaskGrpcInterceptor(unittest.TestCase): |
43 | 59 | """Tests for the DTSDefaultClientInterceptorImpl class.""" |
44 | 60 |
|
@@ -140,34 +156,76 @@ def test_worker_includes_workerid_header(self): |
140 | 156 | self.assertIn("workerid", metadata) |
141 | 157 | self.assertTrue(metadata["workerid"]) |
142 | 158 |
|
| 159 | + def test_client_construction_does_not_acquire_token(self): |
| 160 | + """Token acquisition is deferred from construction to the first request.""" |
| 161 | + credential = _TestTokenCredential() |
143 | 162 |
|
144 | | -class _TestTokenCredential: |
145 | | - def __init__(self): |
146 | | - self._lock = threading.Lock() |
147 | | - self.calls = 0 |
| 163 | + task_hub_client = DurableTaskSchedulerClient( |
| 164 | + host_address=self.server_address, |
| 165 | + secure_channel=False, |
| 166 | + taskhub="test-taskhub", |
| 167 | + token_credential=credential, |
| 168 | + ) |
148 | 169 |
|
149 | | - def get_token(self, _scope): |
150 | | - with self._lock: |
151 | | - self.calls += 1 |
152 | | - call_number = self.calls |
153 | | - time.sleep(0.02) |
154 | | - return AccessToken(f"token-{call_number}", int(time.time()) + 3600) |
| 170 | + self.assertEqual(0, credential.calls, "Constructing a client must not acquire a token") |
| 171 | + |
| 172 | + task_hub_client.get_orchestration_state("test-instance-id") |
| 173 | + |
| 174 | + self.assertEqual(1, credential.calls, "The first request should acquire exactly one token") |
| 175 | + self.assertIn("authorization", self.mock_servicer.captured_metadata) |
| 176 | + |
| 177 | + task_hub_client.get_orchestration_state("test-instance-id") |
| 178 | + |
| 179 | + self.assertEqual(1, credential.calls, "A cached, unexpired token should be reused") |
| 180 | + self.assertEqual(2, self.mock_servicer.requests_received) |
| 181 | + |
| 182 | + def test_worker_construction_does_not_acquire_token(self): |
| 183 | + """Token acquisition is deferred from worker construction to the first request.""" |
| 184 | + credential = _TestTokenCredential() |
| 185 | + |
| 186 | + DurableTaskSchedulerWorker( |
| 187 | + host_address=self.server_address, |
| 188 | + secure_channel=False, |
| 189 | + taskhub="test-taskhub", |
| 190 | + token_credential=credential, |
| 191 | + ) |
| 192 | + |
| 193 | + self.assertEqual(0, credential.calls, "Constructing a worker must not acquire a token") |
155 | 194 |
|
156 | 195 |
|
157 | 196 | class TestAccessTokenManagerThreadSafety(unittest.TestCase): |
| 197 | + |
| 198 | + @staticmethod |
| 199 | + def _get_access_token_concurrently(manager: AccessTokenManager, thread_count: int = 8) -> None: |
| 200 | + threads = [threading.Thread(target=manager.get_access_token) for _ in range(thread_count)] |
| 201 | + for thread in threads: |
| 202 | + thread.start() |
| 203 | + for thread in threads: |
| 204 | + thread.join() |
| 205 | + |
| 206 | + def test_deferred_first_acquisition_performs_single_acquisition(self): |
| 207 | + credential = _TestTokenCredential() |
| 208 | + manager = AccessTokenManager(credential) |
| 209 | + |
| 210 | + self.assertEqual(0, credential.calls, "Constructing the manager must not acquire a token") |
| 211 | + |
| 212 | + self._get_access_token_concurrently(manager) |
| 213 | + |
| 214 | + self.assertEqual(1, credential.calls) |
| 215 | + token = manager.get_access_token() |
| 216 | + self.assertIsNotNone(token) |
| 217 | + self.assertEqual("token-1", token.token) # type: ignore[union-attr] |
| 218 | + self.assertEqual(1, credential.calls, "A cached, unexpired token should be reused") |
| 219 | + |
158 | 220 | def test_concurrent_refresh_performs_single_refresh(self): |
159 | 221 | credential = _TestTokenCredential() |
160 | 222 | manager = AccessTokenManager(credential) |
161 | | - manager.expiry_time = datetime.now(timezone.utc) - timedelta(seconds=1) |
| 223 | + manager.get_access_token() |
| 224 | + self.assertEqual(1, credential.calls) |
162 | 225 |
|
163 | | - threads = [] |
164 | | - for _ in range(8): |
165 | | - thread = threading.Thread(target=manager.get_access_token) |
166 | | - thread.start() |
167 | | - threads.append(thread) |
| 226 | + manager.expiry_time = datetime.now(timezone.utc) - timedelta(seconds=1) |
168 | 227 |
|
169 | | - for thread in threads: |
170 | | - thread.join() |
| 228 | + self._get_access_token_concurrently(manager) |
171 | 229 |
|
172 | 230 | self.assertEqual(2, credential.calls) |
173 | 231 |
|
|
0 commit comments