diff --git a/src/sentry/preprod/api/models/project_preprod_build_details_models.py b/src/sentry/preprod/api/models/project_preprod_build_details_models.py index 7b11bd251fb9a5..ee105943dab4b9 100644 --- a/src/sentry/preprod/api/models/project_preprod_build_details_models.py +++ b/src/sentry/preprod/api/models/project_preprod_build_details_models.py @@ -52,6 +52,12 @@ class BuildDetailsVcsInfo(BaseModel): pr_number: int | None = None +class SizeInfoSizeMetric(BaseModel): + metrics_artifact_type: PreprodArtifactSizeMetrics.MetricsArtifactType + install_size_bytes: int + download_size_bytes: int + + class SizeInfoPending(BaseModel): state: Literal[PreprodArtifactSizeMetrics.SizeAnalysisState.PENDING] = ( PreprodArtifactSizeMetrics.SizeAnalysisState.PENDING @@ -68,8 +74,11 @@ class SizeInfoCompleted(BaseModel): state: Literal[PreprodArtifactSizeMetrics.SizeAnalysisState.COMPLETED] = ( PreprodArtifactSizeMetrics.SizeAnalysisState.COMPLETED ) + # Deprecated, use size_metrics instead install_size_bytes: int + # Deprecated, use size_metrics instead download_size_bytes: int + size_metrics: list[SizeInfoSizeMetric] class SizeInfoFailed(BaseModel): @@ -106,44 +115,65 @@ def platform_from_artifact_type(artifact_type: PreprodArtifact.ArtifactType) -> raise ValueError(f"Unknown artifact type: {artifact_type}") -def to_size_info(size_metrics: None | PreprodArtifactSizeMetrics) -> None | SizeInfo: - if size_metrics is None: +def to_size_info(size_metrics: list[PreprodArtifactSizeMetrics]) -> None | SizeInfo: + if len(size_metrics) == 0: return None - match size_metrics.state: + + main_metric = next( + ( + metric + for metric in size_metrics + if metric.metrics_artifact_type + == PreprodArtifactSizeMetrics.MetricsArtifactType.MAIN_ARTIFACT + ), + # Fallback to the first metric if no main artifact is found + size_metrics[0], + ) + + match main_metric.state: case PreprodArtifactSizeMetrics.SizeAnalysisState.PENDING: return SizeInfoPending() case PreprodArtifactSizeMetrics.SizeAnalysisState.PROCESSING: return SizeInfoProcessing() case PreprodArtifactSizeMetrics.SizeAnalysisState.COMPLETED: - max_install_size = size_metrics.max_install_size - max_download_size = size_metrics.max_download_size + max_install_size = main_metric.max_install_size + max_download_size = main_metric.max_download_size if max_install_size is None or max_download_size is None: raise ValueError( "COMPLETED state requires both max_install_size and max_download_size" ) + return SizeInfoCompleted( - install_size_bytes=max_install_size, download_size_bytes=max_download_size + install_size_bytes=max_install_size, + download_size_bytes=max_download_size, + size_metrics=[ + SizeInfoSizeMetric( + metrics_artifact_type=metric.metrics_artifact_type, + install_size_bytes=metric.max_install_size, + download_size_bytes=metric.max_download_size, + ) + for metric in size_metrics + ], ) case PreprodArtifactSizeMetrics.SizeAnalysisState.FAILED: - error_code = size_metrics.error_code - error_message = size_metrics.error_message + error_code = main_metric.error_code + error_message = main_metric.error_message if error_code is None or error_message is None: raise ValueError("FAILED state requires both error_code and error_message") return SizeInfoFailed(error_code=error_code, error_message=error_message) case _: - raise ValueError(f"Unknown SizeAnalysisState {size_metrics.state}") + raise ValueError(f"Unknown SizeAnalysisState {main_metric.state}") def transform_preprod_artifact_to_build_details( artifact: PreprodArtifact, ) -> BuildDetailsApiResponse: - size_metrics = PreprodArtifactSizeMetrics.objects.filter( + size_metrics_qs = PreprodArtifactSizeMetrics.objects.filter( preprod_artifact=artifact, - metrics_artifact_type=PreprodArtifactSizeMetrics.MetricsArtifactType.MAIN_ARTIFACT, - ).first() + ) - size_info = to_size_info(size_metrics) + size_info = to_size_info(list(size_metrics_qs)) platform = None # artifact_type can be null before preprocessing has completed diff --git a/tests/sentry/preprod/api/endpoints/size_analysis/test_project_preprod_size_analysis_compare.py b/tests/sentry/preprod/api/endpoints/size_analysis/test_project_preprod_size_analysis_compare.py index f7ca752af9333f..8f0d25d6212d91 100644 --- a/tests/sentry/preprod/api/endpoints/size_analysis/test_project_preprod_size_analysis_compare.py +++ b/tests/sentry/preprod/api/endpoints/size_analysis/test_project_preprod_size_analysis_compare.py @@ -231,6 +231,8 @@ def test_get_comparison_success_with_no_matching_base_metric(self): metrics_artifact_type=PreprodArtifactSizeMetrics.MetricsArtifactType.WATCH_ARTIFACT, identifier="watch", state=PreprodArtifactSizeMetrics.SizeAnalysisState.COMPLETED, + max_install_size=500, + max_download_size=250, ) response = self.get_success_response( @@ -395,6 +397,8 @@ def test_get_comparison_multiple_metrics(self): metrics_artifact_type=PreprodArtifactSizeMetrics.MetricsArtifactType.WATCH_ARTIFACT, identifier="watch", state=PreprodArtifactSizeMetrics.SizeAnalysisState.COMPLETED, + max_install_size=500, + max_download_size=250, ) base_watch_metric = PreprodArtifactSizeMetrics.objects.create( @@ -403,6 +407,8 @@ def test_get_comparison_multiple_metrics(self): metrics_artifact_type=PreprodArtifactSizeMetrics.MetricsArtifactType.WATCH_ARTIFACT, identifier="watch", state=PreprodArtifactSizeMetrics.SizeAnalysisState.COMPLETED, + max_install_size=500, + max_download_size=250, ) # Create comparison for watch metrics diff --git a/tests/sentry/preprod/api/models/test_project_preprod_build_details_models.py b/tests/sentry/preprod/api/models/test_project_preprod_build_details_models.py index bb0588cd0c01ef..334576dcabd3a7 100644 --- a/tests/sentry/preprod/api/models/test_project_preprod_build_details_models.py +++ b/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,12 +169,13 @@ 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, ) @@ -143,7 +183,7 @@ def test_to_size_info_completed_state_missing_size_fields(self): 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]))