Merge 9152b4a into c02eef4

Author: lynnagara

src/sentry/hybridcloud/tasks/deliver_webhooks.py

@@ -608,7 +608,6 @@ def perform_region_request(region: Region, payload: WebhookPayload) -> None:
                 "region",
                 {
                     "name": region.name,
-                    "id": region.category,
                     "address": region.address,
                 },
             )

src/sentry/testutils/region.py

@@ -1,128 +1,78 @@
 from __future__ import annotations
 
-from collections.abc import Collection, Generator, Iterable, Sequence
+from collections.abc import Collection, Generator, Sequence
 from contextlib import contextmanager
-from dataclasses import dataclass
 
 from django.test import override_settings
 
 from sentry.types.region import Locality, Region, RegionDirectory, get_global_directory
 
 
-@dataclass(frozen=True)
-class _TemporaryRegionDirectoryState:
-    regions: frozenset[Region]
-    default_region: Region
-
-
 class TestEnvRegionDirectory(RegionDirectory):
     __test__ = False
 
     def __init__(self, regions: Collection[Region]) -> None:
-        super().__init__(regions, self.localities_from_regions(regions))
-        self._tmp_state = _TemporaryRegionDirectoryState(
-            regions=super().regions, default_region=next(iter(regions))
-        )
+        super().__init__(regions, frozenset())
+        self._apply_regions(regions)
 
-    def localities_from_regions(self, regions: Collection[Region]) -> frozenset[Locality]:
-        return frozenset(
-            Locality(name=cell.name, cells=frozenset([cell.name]), category=cell.category)
-            for cell in regions
+    def _apply_regions(self, regions: Collection[Region]) -> None:
+        localities = frozenset(
+            Locality(name=r.name, cells=frozenset([r.name]), category=r.category, visible=r.visible)
+            for r in regions
         )
+        self._cells = frozenset(regions)
+        self._by_name = {r.name: r for r in self._cells}
+        self._localities = localities
+        self._localities_by_name = {loc.name: loc for loc in localities}
+        self._cell_to_locality = {cell_name: loc for loc in localities for cell_name in loc.cells}
 
     @contextmanager
     def swap_state(
         self,
         regions: Sequence[Region],
         local_region: Region | None = None,
     ) -> Generator[None]:
-        monolith_region = regions[0]
-        new_regions = self.regions if regions is None else frozenset(regions)
-        new_state = _TemporaryRegionDirectoryState(
-            regions=new_regions,
-            default_region=local_region or monolith_region,
+        prev_state = (
+            self._cells,
+            self._by_name,
+            self._localities,
+            self._localities_by_name,
+            self._cell_to_locality,
         )
-
-        old_state = self._tmp_state
         try:
-            self._tmp_state = new_state
-
+            self._apply_regions(regions)
+            monolith_region = regions[0]
             with override_settings(SENTRY_MONOLITH_REGION=monolith_region.name):
                 if local_region:
                     with override_settings(SENTRY_REGION=local_region.name):
                         yield
                 else:
                     yield
         finally:
-            self._tmp_state = old_state
+            (
+                self._cells,
+                self._by_name,
+                self._localities,
+                self._localities_by_name,
+                self._cell_to_locality,
+            ) = prev_state
 
     @contextmanager
     def swap_to_default_region(self) -> Generator[None]:
         """Swap to an arbitrary region when entering region mode."""
-        with override_settings(SENTRY_REGION=self._tmp_state.default_region.name):
+        region = next(iter(self._cells))
+        with override_settings(SENTRY_REGION=region.name):
             yield
 
     @contextmanager
     def swap_to_region_by_name(self, region_name: str) -> Generator[None]:
         """Swap to the specified region when entering region mode."""
-
         region = self.get_cell_by_name(region_name)
         if region is None:
             raise Exception("specified swap region not found")
         with override_settings(SENTRY_REGION=region.name):
             yield
 
-    @property
-    def regions(self) -> frozenset[Region]:
-        return self._tmp_state.regions
-
-    @property
-    def localities(self) -> frozenset[Locality]:
-        return frozenset(
-            Locality(
-                name=r.name,
-                cells=frozenset([r.name]),
-                category=r.category,
-                visible=r.visible,
-            )
-            for r in self._tmp_state.regions
-        )
-
-    def get_cell_by_name(self, region_name: str) -> Region | None:
-        match = (r for r in self._tmp_state.regions if r.name == region_name)
-        try:
-            return next(match)
-        except StopIteration:
-            return None
-
-    def get_locality_by_name(self, locality_name: str) -> Locality | None:
-        region = self.get_cell_by_name(locality_name)
-        if region is None:
-            return None
-        return Locality(
-            name=locality_name,
-            cells=frozenset([locality_name]),
-            category=region.category,
-            visible=region.visible,
-        )
-
-    def get_locality_for_cell(self, cell_name: str) -> Locality | None:
-        region = self.get_cell_by_name(cell_name)
-        if region is None:
-            return None
-        return Locality(
-            name=cell_name,
-            cells=frozenset([cell_name]),
-            category=region.category,
-            visible=region.visible,
-        )
-
-    def get_cells_for_locality(self, locality_name: str) -> Iterable[Region]:
-        region = self.get_cell_by_name(locality_name)
-        if region is None:
-            return ()
-        return (region,)
-
 
 def get_test_env_directory() -> TestEnvRegionDirectory:
     directory = get_global_directory()

src/sentry/types/region.py

@@ -167,10 +167,17 @@ def get_locality_by_name(self, locality_name: str) -> Locality | None:
         return self._localities_by_name.get(locality_name)
 
     def get_cells(self, category: RegionCategory | None = None) -> Iterable[Region]:
-        return (r for r in self.regions if (category is None or r.category == category))
+        if category is None:
+            return iter(self.regions)
+
+        return (
+            r
+            for r in self.regions
+            if (loc := self._cell_to_locality.get(r.name)) is not None and loc.category == category
+        )
 
     def get_cell_names(self, category: RegionCategory | None = None) -> Iterable[str]:
-        return (r.name for r in self.regions if (category is None or r.category == category))
+        return (r.name for r in self.get_cells(category))
 
     def get_locality_for_cell(self, cell_name: str) -> Locality | None:
         return self._cell_to_locality.get(cell_name)