|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
3 | 3 | from typing import Any |
| 4 | +from unittest.mock import patch |
4 | 5 |
|
5 | 6 | import orjson |
6 | 7 | from django.test import override_settings |
|
10 | 11 | from sentry.hybridcloud.rpc.service import generate_request_signature |
11 | 12 | from sentry.organizations.services.organization import RpcUserOrganizationContext |
12 | 13 | from sentry.testutils.cases import APITestCase |
| 14 | +from sentry.viewer_context import ActorType, ViewerContext, get_viewer_context |
13 | 15 |
|
14 | 16 |
|
15 | 17 | @override_settings(RPC_SHARED_SECRET=["a-long-value-that-is-hard-to-guess"]) |
@@ -142,3 +144,103 @@ def test_with_invalid_arguments(self) -> None: |
142 | 144 | assert response.data == { |
143 | 145 | "detail": ErrorDetail(string="Malformed request.", code="parse_error") |
144 | 146 | } |
| 147 | + |
| 148 | + def test_viewer_context_propagated_from_meta(self) -> None: |
| 149 | + """ViewerContext in meta is set as the contextvar during dispatch.""" |
| 150 | + organization = self.create_organization() |
| 151 | + captured_contexts: list[ViewerContext | None] = [] |
| 152 | + |
| 153 | + original_dispatch = __import__( |
| 154 | + "sentry.hybridcloud.rpc.service", fromlist=["dispatch_to_local_service"] |
| 155 | + ).dispatch_to_local_service |
| 156 | + |
| 157 | + def capturing_dispatch(*args, **kwargs): |
| 158 | + captured_contexts.append(get_viewer_context()) |
| 159 | + return original_dispatch(*args, **kwargs) |
| 160 | + |
| 161 | + path = self._get_path("organization", "get_organization_by_id") |
| 162 | + data = { |
| 163 | + "args": {"id": organization.id}, |
| 164 | + "meta": { |
| 165 | + "viewer_context": { |
| 166 | + "organization_id": organization.id, |
| 167 | + "user_id": 42, |
| 168 | + "actor_type": "user", |
| 169 | + } |
| 170 | + }, |
| 171 | + } |
| 172 | + |
| 173 | + with patch( |
| 174 | + "sentry.api.endpoints.internal.rpc.dispatch_to_local_service", |
| 175 | + side_effect=capturing_dispatch, |
| 176 | + ): |
| 177 | + response = self._send_post_request(path, data) |
| 178 | + |
| 179 | + assert response.status_code == 200 |
| 180 | + assert len(captured_contexts) == 1 |
| 181 | + ctx = captured_contexts[0] |
| 182 | + assert ctx is not None |
| 183 | + assert ctx.organization_id == organization.id |
| 184 | + assert ctx.user_id == 42 |
| 185 | + assert ctx.actor_type == ActorType.USER |
| 186 | + |
| 187 | + def test_viewer_context_none_when_meta_empty(self) -> None: |
| 188 | + """No ViewerContext set when meta is empty.""" |
| 189 | + organization = self.create_organization() |
| 190 | + captured_contexts: list[ViewerContext | None] = [] |
| 191 | + |
| 192 | + original_dispatch = __import__( |
| 193 | + "sentry.hybridcloud.rpc.service", fromlist=["dispatch_to_local_service"] |
| 194 | + ).dispatch_to_local_service |
| 195 | + |
| 196 | + def capturing_dispatch(*args, **kwargs): |
| 197 | + captured_contexts.append(get_viewer_context()) |
| 198 | + return original_dispatch(*args, **kwargs) |
| 199 | + |
| 200 | + path = self._get_path("organization", "get_organization_by_id") |
| 201 | + data = {"args": {"id": organization.id}, "meta": {}} |
| 202 | + |
| 203 | + with patch( |
| 204 | + "sentry.api.endpoints.internal.rpc.dispatch_to_local_service", |
| 205 | + side_effect=capturing_dispatch, |
| 206 | + ): |
| 207 | + response = self._send_post_request(path, data) |
| 208 | + |
| 209 | + assert response.status_code == 200 |
| 210 | + assert len(captured_contexts) == 1 |
| 211 | + assert captured_contexts[0] is None |
| 212 | + |
| 213 | + def test_viewer_context_roundtrip_through_meta(self) -> None: |
| 214 | + """ViewerContext set on the sending side arrives on the receiving side.""" |
| 215 | + organization = self.create_organization() |
| 216 | + captured_contexts: list[ViewerContext | None] = [] |
| 217 | + |
| 218 | + original_dispatch = __import__( |
| 219 | + "sentry.hybridcloud.rpc.service", fromlist=["dispatch_to_local_service"] |
| 220 | + ).dispatch_to_local_service |
| 221 | + |
| 222 | + def capturing_dispatch(*args, **kwargs): |
| 223 | + captured_contexts.append(get_viewer_context()) |
| 224 | + return original_dispatch(*args, **kwargs) |
| 225 | + |
| 226 | + # Simulate what _send_to_remote_silo builds when ViewerContext is set |
| 227 | + ctx = ViewerContext(organization_id=organization.id, user_id=42, actor_type=ActorType.USER) |
| 228 | + path = self._get_path("organization", "get_organization_by_id") |
| 229 | + data = { |
| 230 | + "args": {"id": organization.id}, |
| 231 | + "meta": {"viewer_context": ctx.serialize()}, |
| 232 | + } |
| 233 | + |
| 234 | + with patch( |
| 235 | + "sentry.api.endpoints.internal.rpc.dispatch_to_local_service", |
| 236 | + side_effect=capturing_dispatch, |
| 237 | + ): |
| 238 | + response = self._send_post_request(path, data) |
| 239 | + |
| 240 | + assert response.status_code == 200 |
| 241 | + assert len(captured_contexts) == 1 |
| 242 | + restored = captured_contexts[0] |
| 243 | + assert restored is not None |
| 244 | + assert restored.organization_id == ctx.organization_id |
| 245 | + assert restored.user_id == ctx.user_id |
| 246 | + assert restored.actor_type == ctx.actor_type |
0 commit comments