1818class TestToSizeInfo (TestCase ):
1919 def test_to_size_info_none_input (self ):
2020 """Test to_size_info returns None when given None input."""
21- result = to_size_info (None )
21+ result = to_size_info ([] )
2222 assert result is None
2323
2424 def test_to_size_info_pending_state (self ):
@@ -27,7 +27,7 @@ def test_to_size_info_pending_state(self):
2727 state = PreprodArtifactSizeMetrics .SizeAnalysisState .PENDING
2828 )
2929
30- result = to_size_info (size_metrics )
30+ result = to_size_info (list ([ size_metrics ]) )
3131
3232 assert isinstance (result , SizeInfoPending )
3333 assert result .state == PreprodArtifactSizeMetrics .SizeAnalysisState .PENDING
@@ -38,7 +38,7 @@ def test_to_size_info_processing_state(self):
3838 state = PreprodArtifactSizeMetrics .SizeAnalysisState .PROCESSING
3939 )
4040
41- result = to_size_info (size_metrics )
41+ result = to_size_info (list ([ size_metrics ]) )
4242
4343 assert isinstance (result , SizeInfoProcessing )
4444 assert result .state == PreprodArtifactSizeMetrics .SizeAnalysisState .PROCESSING
@@ -47,17 +47,54 @@ def test_to_size_info_completed_state(self):
4747 """Test to_size_info returns SizeInfoCompleted for COMPLETED state."""
4848 size_metrics = PreprodArtifactSizeMetrics (
4949 state = PreprodArtifactSizeMetrics .SizeAnalysisState .COMPLETED ,
50+ metrics_artifact_type = PreprodArtifactSizeMetrics .MetricsArtifactType .MAIN_ARTIFACT ,
5051 max_install_size = 1024000 ,
5152 max_download_size = 512000 ,
5253 )
5354
54- result = to_size_info (size_metrics )
55+ result = to_size_info (list ([ size_metrics ]) )
5556
5657 assert isinstance (result , SizeInfoCompleted )
5758 assert result .state == PreprodArtifactSizeMetrics .SizeAnalysisState .COMPLETED
5859 assert result .install_size_bytes == 1024000
5960 assert result .download_size_bytes == 512000
6061
62+ def test_to_size_info_completed_state_with_multiple_metrics (self ):
63+ """Test to_size_info returns SizeInfoCompleted for COMPLETED state with multiple metrics."""
64+ size_metrics = [
65+ PreprodArtifactSizeMetrics (
66+ state = PreprodArtifactSizeMetrics .SizeAnalysisState .COMPLETED ,
67+ metrics_artifact_type = PreprodArtifactSizeMetrics .MetricsArtifactType .MAIN_ARTIFACT ,
68+ max_install_size = 1024000 ,
69+ max_download_size = 512000 ,
70+ ),
71+ PreprodArtifactSizeMetrics (
72+ state = PreprodArtifactSizeMetrics .SizeAnalysisState .COMPLETED ,
73+ metrics_artifact_type = PreprodArtifactSizeMetrics .MetricsArtifactType .WATCH_ARTIFACT ,
74+ max_install_size = 512000 ,
75+ max_download_size = 256000 ,
76+ ),
77+ ]
78+
79+ result = to_size_info (size_metrics )
80+
81+ assert isinstance (result , SizeInfoCompleted )
82+ assert result .install_size_bytes == 1024000
83+ assert result .download_size_bytes == 512000
84+ assert len (result .size_metrics ) == 2
85+ assert (
86+ result .size_metrics [0 ].metrics_artifact_type
87+ == PreprodArtifactSizeMetrics .MetricsArtifactType .MAIN_ARTIFACT
88+ )
89+ assert result .size_metrics [0 ].install_size_bytes == 1024000
90+ assert result .size_metrics [0 ].download_size_bytes == 512000
91+ assert (
92+ result .size_metrics [1 ].metrics_artifact_type
93+ == PreprodArtifactSizeMetrics .MetricsArtifactType .WATCH_ARTIFACT
94+ )
95+ assert result .size_metrics [1 ].install_size_bytes == 512000
96+ assert result .size_metrics [1 ].download_size_bytes == 256000
97+
6198 def test_to_size_info_failed_state (self ):
6299 """Test to_size_info returns SizeInfoFailed for FAILED state."""
63100 size_metrics = PreprodArtifactSizeMetrics (
@@ -66,7 +103,7 @@ def test_to_size_info_failed_state(self):
66103 error_message = "Analysis timed out after 30 minutes" ,
67104 )
68105
69- result = to_size_info (size_metrics )
106+ result = to_size_info (list ([ size_metrics ]) )
70107
71108 assert isinstance (result , SizeInfoFailed )
72109 assert result .state == PreprodArtifactSizeMetrics .SizeAnalysisState .FAILED
@@ -91,7 +128,7 @@ def test_to_size_info_failed_state_with_different_error_codes(self):
91128 error_message = error_message ,
92129 )
93130
94- result = to_size_info (size_metrics )
131+ result = to_size_info (list ([ size_metrics ]) )
95132
96133 assert isinstance (result , SizeInfoFailed )
97134 assert result .error_code == error_code
@@ -101,11 +138,12 @@ def test_to_size_info_completed_with_zero_sizes(self):
101138 """Test to_size_info handles completed state with zero sizes."""
102139 size_metrics = PreprodArtifactSizeMetrics (
103140 state = PreprodArtifactSizeMetrics .SizeAnalysisState .COMPLETED ,
141+ metrics_artifact_type = PreprodArtifactSizeMetrics .MetricsArtifactType .MAIN_ARTIFACT ,
104142 max_install_size = 0 ,
105143 max_download_size = 0 ,
106144 )
107145
108- result = to_size_info (size_metrics )
146+ result = to_size_info (list ([ size_metrics ]) )
109147
110148 assert isinstance (result , SizeInfoCompleted )
111149 assert result .install_size_bytes == 0
@@ -115,11 +153,12 @@ def test_to_size_info_completed_with_large_sizes(self):
115153 """Test to_size_info handles completed state with large file sizes."""
116154 size_metrics = PreprodArtifactSizeMetrics (
117155 state = PreprodArtifactSizeMetrics .SizeAnalysisState .COMPLETED ,
156+ metrics_artifact_type = PreprodArtifactSizeMetrics .MetricsArtifactType .MAIN_ARTIFACT ,
118157 max_install_size = 5000000000 , # ~5GB
119158 max_download_size = 2000000000 , # ~2GB
120159 )
121160
122- result = to_size_info (size_metrics )
161+ result = to_size_info (list ([ size_metrics ]) )
123162
124163 assert isinstance (result , SizeInfoCompleted )
125164 assert result .install_size_bytes == 5000000000
@@ -130,20 +169,21 @@ def test_to_size_info_invalid_state_raises_error(self):
130169 size_metrics = PreprodArtifactSizeMetrics (state = 999 ) # Invalid state
131170
132171 with pytest .raises (ValueError , match = "Unknown SizeAnalysisState 999" ):
133- to_size_info (size_metrics )
172+ to_size_info (list ([ size_metrics ]) )
134173
135174 def test_to_size_info_completed_state_missing_size_fields (self ):
136175 """Test to_size_info raises ValueError when COMPLETED state has None size fields."""
137176 size_metrics = PreprodArtifactSizeMetrics (
138177 state = PreprodArtifactSizeMetrics .SizeAnalysisState .COMPLETED ,
178+ metrics_artifact_type = PreprodArtifactSizeMetrics .MetricsArtifactType .MAIN_ARTIFACT ,
139179 max_install_size = None ,
140180 max_download_size = None ,
141181 )
142182
143183 with pytest .raises (
144184 ValueError , match = "COMPLETED state requires both max_install_size and max_download_size"
145185 ):
146- to_size_info (size_metrics )
186+ to_size_info (list ([ size_metrics ]) )
147187
148188 def test_to_size_info_failed_state_no_error_code (self ):
149189 """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):
156196 with pytest .raises (
157197 ValueError , match = "FAILED state requires both error_code and error_message"
158198 ):
159- to_size_info (size_metrics )
199+ to_size_info (list ([ size_metrics ]) )
160200
161201 def test_to_size_info_failed_state_no_error_message (self ):
162202 """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):
169209 with pytest .raises (
170210 ValueError , match = "FAILED state requires both error_code and error_message"
171211 ):
172- to_size_info (size_metrics )
212+ to_size_info (list ([ size_metrics ]) )
0 commit comments