chore: Format

Author: anth-volk

projects/policyengine-api-tagger/src/policyengine_api_tagger/api/revision_cleanup.py

@@ -20,6 +20,7 @@
 
 class TagInfo(BaseModel):
     """Information about a traffic tag."""
+
     tag: str
     revision: str
     country: str  # "us" or "uk"
@@ -77,7 +78,7 @@ def _parse_tag(self, tag: str, revision: str) -> TagInfo | None:
         Example: country-us-model-1-459-0 -> country=us, version=(1, 459, 0)
         """
         # Match pattern: country-{us|uk}-model-{version}
-        match = re.match(r'^country-(us|uk)-model-(.+)$', tag)
+        match = re.match(r"^country-(us|uk)-model-(.+)$", tag)
         if not match:
             log.debug(f"Tag '{tag}' doesn't match expected pattern, skipping")
             return None
@@ -86,11 +87,11 @@ def _parse_tag(self, tag: str, revision: str) -> TagInfo | None:
         version_with_dashes = match.group(2)
 
         # Convert dashes back to dots for the version string
-        version_str = version_with_dashes.replace('-', '.')
+        version_str = version_with_dashes.replace("-", ".")
 
         # Parse version into tuple of integers for comparison
         try:
-            version_parts = tuple(int(p) for p in version_str.split('.'))
+            version_parts = tuple(int(p) for p in version_str.split("."))
         except ValueError:
             log.warning(f"Could not parse version from tag '{tag}': {version_str}")
             return None
@@ -109,7 +110,9 @@ async def _get_service(self) -> Service:
         service_name = self._get_service_name()
         return await client.get_service(name=service_name)
 
-    async def _update_service_traffic(self, service: Service, tags_to_keep: list[TagInfo]) -> None:
+    async def _update_service_traffic(
+        self, service: Service, tags_to_keep: list[TagInfo]
+    ) -> None:
         """
         Update the service traffic configuration to only include specified tags.
 
@@ -130,19 +133,24 @@ async def _update_service_traffic(self, service: Service, tags_to_keep: list[Tag
         # Add the tags we want to keep (with percent=0)
         for tag_info in tags_to_keep:
             from google.cloud.run_v2 import TrafficTarget
-            new_traffic.append(TrafficTarget(
-                percent=0,
-                revision=tag_info.revision,
-                tag=tag_info.tag,
-            ))
+
+            new_traffic.append(
+                TrafficTarget(
+                    percent=0,
+                    revision=tag_info.revision,
+                    tag=tag_info.tag,
+                )
+            )
 
         # Update the service
         service.traffic = new_traffic
 
         request = UpdateServiceRequest(service=service)
         await client.update_service(request=request)
 
-    async def _analyze_tags(self, keep_count: int) -> tuple[
+    async def _analyze_tags(
+        self, keep_count: int
+    ) -> tuple[
         Service | None,
         list[TagInfo],
         list[TagInfo],
@@ -226,7 +234,15 @@ async def _analyze_tags(self, keep_count: int) -> tuple[
 
         log.info(f"Keeping {len(tags_to_keep)} tags, removing {len(tags_removed)} tags")
 
-        return service, all_tags, tags_to_keep, newest_us, newest_uk, tags_removed, errors
+        return (
+            service,
+            all_tags,
+            tags_to_keep,
+            newest_us,
+            newest_uk,
+            tags_removed,
+            errors,
+        )
 
     async def preview(self, keep_count: int = 40) -> CleanupResult:
         """
@@ -238,8 +254,15 @@ async def preview(self, keep_count: int = 40) -> CleanupResult:
         Returns:
             CleanupResult showing what would be kept/removed
         """
-        service, all_tags, tags_to_keep, newest_us, newest_uk, tags_removed, errors = \
-            await self._analyze_tags(keep_count)
+        (
+            service,
+            all_tags,
+            tags_to_keep,
+            newest_us,
+            newest_uk,
+            tags_removed,
+            errors,
+        ) = await self._analyze_tags(keep_count)
 
         if service is None:
             return CleanupResult(
@@ -276,8 +299,15 @@ async def cleanup(self, keep_count: int = 40) -> CleanupResult:
         Returns:
             CleanupResult with details of what was cleaned up
         """
-        service, all_tags, tags_to_keep, newest_us, newest_uk, tags_removed, errors = \
-            await self._analyze_tags(keep_count)
+        (
+            service,
+            all_tags,
+            tags_to_keep,
+            newest_us,
+            newest_uk,
+            tags_removed,
+            errors,
+        ) = await self._analyze_tags(keep_count)
 
         if service is None:
             return CleanupResult(

projects/policyengine-api-tagger/src/policyengine_api_tagger/api/routes.py

@@ -23,7 +23,9 @@ async def get_tag_uri(country: str, model_version: str) -> str:
         return uri
 
     @router.post("/cleanup")
-    async def cleanup_old_revisions(keep: int = 40, dry_run: bool = False) -> CleanupResult:
+    async def cleanup_old_revisions(
+        keep: int = 40, dry_run: bool = False
+    ) -> CleanupResult:
         """
         Clean up old traffic tags, keeping the specified number.
 

projects/policyengine-api-tagger/tests/api/test_cleanup_integration.py

@@ -71,9 +71,7 @@ def client(self, mock_cloudrun_service):
 
             yield TestClient(app)
 
-    def test_cleanup_returns_success_with_tags(
-        self, client, mock_cloudrun_service
-    ):
+    def test_cleanup_returns_success_with_tags(self, client, mock_cloudrun_service):
         """Cleanup succeeds and returns tag information."""
         mock_cloudrun_service["service"].traffic = [
             make_mock_traffic_entry("country-us-model-1-459-0", "rev-us"),
@@ -90,9 +88,7 @@ def test_cleanup_returns_success_with_tags(
         assert result["newest_us_tag"] == "country-us-model-1-459-0"
         assert result["newest_uk_tag"] == "country-uk-model-2-65-9"
 
-    def test_cleanup_dry_run_does_not_modify(
-        self, client, mock_cloudrun_service
-    ):
+    def test_cleanup_dry_run_does_not_modify(self, client, mock_cloudrun_service):
         """Cleanup with dry_run=true should NOT call update_service."""
         mock_cloudrun_service["service"].traffic = [
             make_mock_traffic_entry("country-us-model-1-100-0", "rev-1"),
@@ -112,9 +108,7 @@ def test_cleanup_dry_run_does_not_modify(
         # CRITICAL: update_service should NOT have been called
         mock_cloudrun_service["client"].update_service.assert_not_called()
 
-    def test_cleanup_identifies_safeguards(
-        self, client, mock_cloudrun_service
-    ):
+    def test_cleanup_identifies_safeguards(self, client, mock_cloudrun_service):
         """Cleanup correctly identifies newest US and UK tags."""
         mock_cloudrun_service["service"].traffic = [
             make_mock_traffic_entry("country-us-model-1-100-0", "rev-1"),
@@ -130,9 +124,7 @@ def test_cleanup_identifies_safeguards(
         assert result["newest_us_tag"] == "country-us-model-1-459-0"
         assert result["newest_uk_tag"] == "country-uk-model-2-65-9"
 
-    def test_cleanup_respects_keep_count(
-        self, client, mock_cloudrun_service
-    ):
+    def test_cleanup_respects_keep_count(self, client, mock_cloudrun_service):
         """Cleanup keeps the correct number of tags."""
         mock_cloudrun_service["service"].traffic = [
             make_mock_traffic_entry("country-us-model-1-100-0", "rev-1"),
@@ -156,9 +148,7 @@ def test_cleanup_validates_keep_minimum(self, client, mock_cloudrun_service):
         assert response.status_code == 400
         assert "at least 2" in response.json()["detail"]
 
-    def test_cleanup_handles_no_tags(
-        self, client, mock_cloudrun_service
-    ):
+    def test_cleanup_handles_no_tags(self, client, mock_cloudrun_service):
         """Cleanup handles service with no tags."""
         mock_cloudrun_service["service"].traffic = [
             make_mock_traffic_entry(None, "rev-main", percent=100),

projects/policyengine-api-tagger/tests/api/test_revision_cleanup.py

@@ -101,8 +101,15 @@ async def test_identifies_newest_us_and_uk_tags(self, cleanup):
         with patch.object(cleanup, "_get_service", new_callable=AsyncMock) as mock_get:
             mock_get.return_value = mock_service
 
-            service, all_tags, tags_to_keep, newest_us, newest_uk, tags_removed, errors = \
-                await cleanup._analyze_tags(keep_count=40)
+            (
+                service,
+                all_tags,
+                tags_to_keep,
+                newest_us,
+                newest_uk,
+                tags_removed,
+                errors,
+            ) = await cleanup._analyze_tags(keep_count=40)
 
         assert newest_us is not None
         assert newest_us.tag == "country-us-model-1-459-0"
@@ -126,8 +133,15 @@ async def test_keeps_safeguards_plus_newest(self, cleanup):
         with patch.object(cleanup, "_get_service", new_callable=AsyncMock) as mock_get:
             mock_get.return_value = mock_service
 
-            service, all_tags, tags_to_keep, newest_us, newest_uk, tags_removed, errors = \
-                await cleanup._analyze_tags(keep_count=4)
+            (
+                service,
+                all_tags,
+                tags_to_keep,
+                newest_us,
+                newest_uk,
+                tags_removed,
+                errors,
+            ) = await cleanup._analyze_tags(keep_count=4)
 
         # Should keep: newest US, newest UK, then next 2 newest by version
         assert len(tags_to_keep) == 4
@@ -152,8 +166,15 @@ async def test_handles_keep_count_less_than_2(self, cleanup):
         with patch.object(cleanup, "_get_service", new_callable=AsyncMock) as mock_get:
             mock_get.return_value = mock_service
 
-            service, all_tags, tags_to_keep, newest_us, newest_uk, tags_removed, errors = \
-                await cleanup._analyze_tags(keep_count=1)  # Should become 2
+            (
+                service,
+                all_tags,
+                tags_to_keep,
+                newest_us,
+                newest_uk,
+                tags_removed,
+                errors,
+            ) = await cleanup._analyze_tags(keep_count=1)  # Should become 2
 
         # Should still keep both safeguards
         assert len(tags_to_keep) == 2
@@ -169,8 +190,15 @@ async def test_handles_no_tags(self, cleanup):
         with patch.object(cleanup, "_get_service", new_callable=AsyncMock) as mock_get:
             mock_get.return_value = mock_service
 
-            service, all_tags, tags_to_keep, newest_us, newest_uk, tags_removed, errors = \
-                await cleanup._analyze_tags(keep_count=40)
+            (
+                service,
+                all_tags,
+                tags_to_keep,
+                newest_us,
+                newest_uk,
+                tags_removed,
+                errors,
+            ) = await cleanup._analyze_tags(keep_count=40)
 
         assert len(all_tags) == 0
         assert len(tags_to_keep) == 0
@@ -189,8 +217,15 @@ async def test_handles_only_us_tags(self, cleanup):
         with patch.object(cleanup, "_get_service", new_callable=AsyncMock) as mock_get:
             mock_get.return_value = mock_service
 
-            service, all_tags, tags_to_keep, newest_us, newest_uk, tags_removed, errors = \
-                await cleanup._analyze_tags(keep_count=40)
+            (
+                service,
+                all_tags,
+                tags_to_keep,
+                newest_us,
+                newest_uk,
+                tags_removed,
+                errors,
+            ) = await cleanup._analyze_tags(keep_count=40)
 
         assert newest_us is not None
         assert newest_us.tag == "country-us-model-1-200-0"
@@ -202,8 +237,15 @@ async def test_handles_service_error(self, cleanup):
         with patch.object(cleanup, "_get_service", new_callable=AsyncMock) as mock_get:
             mock_get.side_effect = Exception("Cloud Run API error")
 
-            service, all_tags, tags_to_keep, newest_us, newest_uk, tags_removed, errors = \
-                await cleanup._analyze_tags(keep_count=40)
+            (
+                service,
+                all_tags,
+                tags_to_keep,
+                newest_us,
+                newest_uk,
+                tags_removed,
+                errors,
+            ) = await cleanup._analyze_tags(keep_count=40)
 
         assert service is None
         assert len(errors) == 1
@@ -250,7 +292,9 @@ async def test_preview_does_not_call_update(self, cleanup):
 
         with (
             patch.object(cleanup, "_get_service", new_callable=AsyncMock) as mock_get,
-            patch.object(cleanup, "_update_service_traffic", new_callable=AsyncMock) as mock_update,
+            patch.object(
+                cleanup, "_update_service_traffic", new_callable=AsyncMock
+            ) as mock_update,
         ):
             mock_get.return_value = mock_service
 
@@ -278,7 +322,9 @@ async def test_calls_update_when_tags_to_remove(self, cleanup):
 
         with (
             patch.object(cleanup, "_get_service", new_callable=AsyncMock) as mock_get,
-            patch.object(cleanup, "_update_service_traffic", new_callable=AsyncMock) as mock_update,
+            patch.object(
+                cleanup, "_update_service_traffic", new_callable=AsyncMock
+            ) as mock_update,
         ):
             mock_get.return_value = mock_service
 
@@ -299,7 +345,9 @@ async def test_does_not_call_update_when_nothing_to_remove(self, cleanup):
 
         with (
             patch.object(cleanup, "_get_service", new_callable=AsyncMock) as mock_get,
-            patch.object(cleanup, "_update_service_traffic", new_callable=AsyncMock) as mock_update,
+            patch.object(
+                cleanup, "_update_service_traffic", new_callable=AsyncMock
+            ) as mock_update,
         ):
             mock_get.return_value = mock_service
 
@@ -322,7 +370,9 @@ async def test_handles_update_failure(self, cleanup):
 
         with (
             patch.object(cleanup, "_get_service", new_callable=AsyncMock) as mock_get,
-            patch.object(cleanup, "_update_service_traffic", new_callable=AsyncMock) as mock_update,
+            patch.object(
+                cleanup, "_update_service_traffic", new_callable=AsyncMock
+            ) as mock_update,
         ):
             mock_get.return_value = mock_service
             mock_update.side_effect = Exception("Update failed")
@@ -354,7 +404,9 @@ async def capture_update(service, tags_to_keep):
 
         with (
             patch.object(cleanup, "_get_service", new_callable=AsyncMock) as mock_get,
-            patch.object(cleanup, "_update_service_traffic", side_effect=capture_update) as mock_update,
+            patch.object(
+                cleanup, "_update_service_traffic", side_effect=capture_update
+            ) as mock_update,
         ):
             mock_get.return_value = mock_service
 
@@ -388,8 +440,15 @@ async def test_sorts_versions_correctly(self, cleanup):
         with patch.object(cleanup, "_get_service", new_callable=AsyncMock) as mock_get:
             mock_get.return_value = mock_service
 
-            service, all_tags, tags_to_keep, newest_us, newest_uk, tags_removed, errors = \
-                await cleanup._analyze_tags(keep_count=2)
+            (
+                service,
+                all_tags,
+                tags_to_keep,
+                newest_us,
+                newest_uk,
+                tags_removed,
+                errors,
+            ) = await cleanup._analyze_tags(keep_count=2)
 
         # 1.100.0 > 1.10.0 > 1.9.0 numerically
         assert newest_us.tag == "country-us-model-1-100-0"