Fix tests

Author: rbro112

src/sentry/preprod/api/models/project_preprod_build_details_models.py

@@ -120,10 +120,13 @@ def to_size_info(size_metrics: list[PreprodArtifactSizeMetrics]) -> None | SizeI
         return None
 
     main_metric = next(
-        metric
-        for metric in size_metrics
-        if metric.metrics_artifact_type
-        == PreprodArtifactSizeMetrics.MetricsArtifactType.MAIN_ARTIFACT
+        (
+            metric
+            for metric in size_metrics
+            if metric.metrics_artifact_type
+            == PreprodArtifactSizeMetrics.MetricsArtifactType.MAIN_ARTIFACT
+        ),
+        size_metrics[0],
     )
 
     match main_metric.state:

tests/sentry/preprod/api/models/test_project_preprod_build_details_models.py

@@ -18,7 +18,7 @@
 class TestToSizeInfo(TestCase):
     def test_to_size_info_none_input(self):
         """Test to_size_info returns None when given None input."""
-        result = to_size_info(None)
+        result = to_size_info([])
         assert result is None
 
     def test_to_size_info_pending_state(self):
@@ -27,7 +27,7 @@ def test_to_size_info_pending_state(self):
             state=PreprodArtifactSizeMetrics.SizeAnalysisState.PENDING
         )
 
-        result = to_size_info(size_metrics)
+        result = to_size_info(list([size_metrics]))
 
         assert isinstance(result, SizeInfoPending)
         assert result.state == PreprodArtifactSizeMetrics.SizeAnalysisState.PENDING
@@ -38,7 +38,7 @@ def test_to_size_info_processing_state(self):
             state=PreprodArtifactSizeMetrics.SizeAnalysisState.PROCESSING
         )
 
-        result = to_size_info(size_metrics)
+        result = to_size_info(list([size_metrics]))
 
         assert isinstance(result, SizeInfoProcessing)
         assert result.state == PreprodArtifactSizeMetrics.SizeAnalysisState.PROCESSING
@@ -47,17 +47,54 @@ def test_to_size_info_completed_state(self):
         """Test to_size_info returns SizeInfoCompleted for COMPLETED state."""
         size_metrics = PreprodArtifactSizeMetrics(
             state=PreprodArtifactSizeMetrics.SizeAnalysisState.COMPLETED,
+            metrics_artifact_type=PreprodArtifactSizeMetrics.MetricsArtifactType.MAIN_ARTIFACT,
             max_install_size=1024000,
             max_download_size=512000,
         )
 
-        result = to_size_info(size_metrics)
+        result = to_size_info(list([size_metrics]))
 
         assert isinstance(result, SizeInfoCompleted)
         assert result.state == PreprodArtifactSizeMetrics.SizeAnalysisState.COMPLETED
         assert result.install_size_bytes == 1024000
         assert result.download_size_bytes == 512000
 
+    def test_to_size_info_completed_state_with_multiple_metrics(self):
+        """Test to_size_info returns SizeInfoCompleted for COMPLETED state with multiple metrics."""
+        size_metrics = [
+            PreprodArtifactSizeMetrics(
+                state=PreprodArtifactSizeMetrics.SizeAnalysisState.COMPLETED,
+                metrics_artifact_type=PreprodArtifactSizeMetrics.MetricsArtifactType.MAIN_ARTIFACT,
+                max_install_size=1024000,
+                max_download_size=512000,
+            ),
+            PreprodArtifactSizeMetrics(
+                state=PreprodArtifactSizeMetrics.SizeAnalysisState.COMPLETED,
+                metrics_artifact_type=PreprodArtifactSizeMetrics.MetricsArtifactType.WATCH_ARTIFACT,
+                max_install_size=512000,
+                max_download_size=256000,
+            ),
+        ]
+
+        result = to_size_info(size_metrics)
+
+        assert isinstance(result, SizeInfoCompleted)
+        assert result.install_size_bytes == 1024000
+        assert result.download_size_bytes == 512000
+        assert len(result.size_metrics) == 2
+        assert (
+            result.size_metrics[0].metrics_artifact_type
+            == PreprodArtifactSizeMetrics.MetricsArtifactType.MAIN_ARTIFACT
+        )
+        assert result.size_metrics[0].install_size_bytes == 1024000
+        assert result.size_metrics[0].download_size_bytes == 512000
+        assert (
+            result.size_metrics[1].metrics_artifact_type
+            == PreprodArtifactSizeMetrics.MetricsArtifactType.WATCH_ARTIFACT
+        )
+        assert result.size_metrics[1].install_size_bytes == 512000
+        assert result.size_metrics[1].download_size_bytes == 256000
+
     def test_to_size_info_failed_state(self):
         """Test to_size_info returns SizeInfoFailed for FAILED state."""
         size_metrics = PreprodArtifactSizeMetrics(
@@ -66,7 +103,7 @@ def test_to_size_info_failed_state(self):
             error_message="Analysis timed out after 30 minutes",
         )
 
-        result = to_size_info(size_metrics)
+        result = to_size_info(list([size_metrics]))
 
         assert isinstance(result, SizeInfoFailed)
         assert result.state == PreprodArtifactSizeMetrics.SizeAnalysisState.FAILED
@@ -91,7 +128,7 @@ def test_to_size_info_failed_state_with_different_error_codes(self):
                 error_message=error_message,
             )
 
-            result = to_size_info(size_metrics)
+            result = to_size_info(list([size_metrics]))
 
             assert isinstance(result, SizeInfoFailed)
             assert result.error_code == error_code
@@ -101,11 +138,12 @@ def test_to_size_info_completed_with_zero_sizes(self):
         """Test to_size_info handles completed state with zero sizes."""
         size_metrics = PreprodArtifactSizeMetrics(
             state=PreprodArtifactSizeMetrics.SizeAnalysisState.COMPLETED,
+            metrics_artifact_type=PreprodArtifactSizeMetrics.MetricsArtifactType.MAIN_ARTIFACT,
             max_install_size=0,
             max_download_size=0,
         )
 
-        result = to_size_info(size_metrics)
+        result = to_size_info(list([size_metrics]))
 
         assert isinstance(result, SizeInfoCompleted)
         assert result.install_size_bytes == 0
@@ -115,11 +153,12 @@ def test_to_size_info_completed_with_large_sizes(self):
         """Test to_size_info handles completed state with large file sizes."""
         size_metrics = PreprodArtifactSizeMetrics(
             state=PreprodArtifactSizeMetrics.SizeAnalysisState.COMPLETED,
+            metrics_artifact_type=PreprodArtifactSizeMetrics.MetricsArtifactType.MAIN_ARTIFACT,
             max_install_size=5000000000,  # ~5GB
             max_download_size=2000000000,  # ~2GB
         )
 
-        result = to_size_info(size_metrics)
+        result = to_size_info(list([size_metrics]))
 
         assert isinstance(result, SizeInfoCompleted)
         assert result.install_size_bytes == 5000000000
@@ -130,20 +169,21 @@ def test_to_size_info_invalid_state_raises_error(self):
         size_metrics = PreprodArtifactSizeMetrics(state=999)  # Invalid state
 
         with pytest.raises(ValueError, match="Unknown SizeAnalysisState 999"):
-            to_size_info(size_metrics)
+            to_size_info(list([size_metrics]))
 
     def test_to_size_info_completed_state_missing_size_fields(self):
         """Test to_size_info raises ValueError when COMPLETED state has None size fields."""
         size_metrics = PreprodArtifactSizeMetrics(
             state=PreprodArtifactSizeMetrics.SizeAnalysisState.COMPLETED,
+            metrics_artifact_type=PreprodArtifactSizeMetrics.MetricsArtifactType.MAIN_ARTIFACT,
             max_install_size=None,
             max_download_size=None,
         )
 
         with pytest.raises(
             ValueError, match="COMPLETED state requires both max_install_size and max_download_size"
         ):
-            to_size_info(size_metrics)
+            to_size_info(list([size_metrics]))
 
     def test_to_size_info_failed_state_no_error_code(self):
         """Test to_size_info raises ValueError when FAILED state has only error_code."""
@@ -156,7 +196,7 @@ def test_to_size_info_failed_state_no_error_code(self):
         with pytest.raises(
             ValueError, match="FAILED state requires both error_code and error_message"
         ):
-            to_size_info(size_metrics)
+            to_size_info(list([size_metrics]))
 
     def test_to_size_info_failed_state_no_error_message(self):
         """Test to_size_info raises ValueError when FAILED state has only error_message."""
@@ -169,4 +209,4 @@ def test_to_size_info_failed_state_no_error_message(self):
         with pytest.raises(
             ValueError, match="FAILED state requires both error_code and error_message"
         ):
-            to_size_info(size_metrics)
+            to_size_info(list([size_metrics]))