|
18 | 18 | from mcp.types import ServerCapabilities |
19 | 19 |
|
20 | 20 |
|
| 21 | +def _simulate_tenant_binding(session: ServerSession, tenant_id_value: str) -> None: |
| 22 | + """Simulate the set-once tenant binding logic from lowlevel/server.py.""" |
| 23 | + access_token = AccessToken( |
| 24 | + token=f"token-{tenant_id_value}", |
| 25 | + client_id="client", |
| 26 | + scopes=["read"], |
| 27 | + expires_at=int(time.time()) + 3600, |
| 28 | + tenant_id=tenant_id_value, |
| 29 | + ) |
| 30 | + user = AuthenticatedUser(access_token) |
| 31 | + context_token = auth_context_var.set(user) |
| 32 | + try: |
| 33 | + tenant_id = get_tenant_id() |
| 34 | + if tenant_id is not None and session.tenant_id is None: |
| 35 | + session.tenant_id = tenant_id |
| 36 | + finally: |
| 37 | + auth_context_var.reset(context_token) |
| 38 | + |
| 39 | + |
21 | 40 | @pytest.fixture |
22 | 41 | def init_options() -> InitializationOptions: |
23 | 42 | """Create initialization options for testing.""" |
@@ -177,41 +196,12 @@ async def test_session_tenant_id_set_from_auth_context_on_first_request(init_opt |
177 | 196 |
|
178 | 197 | # Simulate what lowlevel/server.py does: set session.tenant_id |
179 | 198 | # from auth context on first request |
180 | | - access_token = AccessToken( |
181 | | - token="token-first", |
182 | | - client_id="client", |
183 | | - scopes=["read"], |
184 | | - expires_at=int(time.time()) + 3600, |
185 | | - tenant_id="tenant-first", |
186 | | - ) |
187 | | - user = AuthenticatedUser(access_token) |
188 | | - context_token = auth_context_var.set(user) |
189 | | - try: |
190 | | - tenant_id = get_tenant_id() |
191 | | - if tenant_id is not None and session.tenant_id is None: |
192 | | - session.tenant_id = tenant_id |
193 | | - finally: |
194 | | - auth_context_var.reset(context_token) |
195 | | - |
| 199 | + _simulate_tenant_binding(session, "tenant-first") |
196 | 200 | assert session.tenant_id == "tenant-first" |
197 | 201 |
|
198 | 202 | # Simulate a second request with a different tenant — |
199 | 203 | # session.tenant_id should NOT change (set-once on first request) |
200 | | - access_token2 = AccessToken( |
201 | | - token="token-second", |
202 | | - client_id="client", |
203 | | - scopes=["read"], |
204 | | - expires_at=int(time.time()) + 3600, |
205 | | - tenant_id="tenant-second", |
206 | | - ) |
207 | | - user2 = AuthenticatedUser(access_token2) |
208 | | - context_token2 = auth_context_var.set(user2) |
209 | | - try: |
210 | | - tenant_id = get_tenant_id() |
211 | | - if tenant_id is not None and session.tenant_id is None: |
212 | | - session.tenant_id = tenant_id |
213 | | - finally: |
214 | | - auth_context_var.reset(context_token2) |
| 204 | + _simulate_tenant_binding(session, "tenant-second") |
215 | 205 |
|
216 | 206 | # Still the first tenant — not overwritten |
217 | 207 | assert session.tenant_id == "tenant-first" |
|
0 commit comments