@@ -128,14 +128,6 @@ def mock_storage_blob_download_to_filename():
128
128
yield mock_blob_download_to_filename
129
129
130
130
131
- @pytest .fixture ()
132
- def mock_bucket_not_exist ():
133
- with patch ("google.cloud.storage.Blob.from_string" ) as mock_bucket_not_exist , patch (
134
- "google.cloud.storage.Bucket.exists" , return_value = False
135
- ):
136
- yield mock_bucket_not_exist
137
-
138
-
139
131
def test_invalid_region_raises_with_invalid_region ():
140
132
with pytest .raises (ValueError ):
141
133
aiplatform .utils .validate_region (region = "us-east5" )
@@ -587,8 +579,13 @@ def test_generate_gcs_directory_for_pipeline_artifacts(self):
587
579
)
588
580
assert output == "gs://project-vertex-pipelines-us-central1/output_artifacts/"
589
581
582
+ @patch .object (storage .Bucket , "exists" , return_value = False )
583
+ @patch .object (storage , "Client" )
584
+ @patch .object (
585
+ gcs_utils .resource_manager_utils , "get_project_number" , return_value = 12345
586
+ )
590
587
def test_create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist (
591
- self , mock_bucket_not_exist , mock_storage_client
588
+ self , mock_bucket_not_exist , mock_storage_client , mock_get_project_number
592
589
):
593
590
output = (
594
591
gcs_utils .create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist (
@@ -597,6 +594,7 @@ def test_create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
597
594
)
598
595
assert mock_storage_client .called
599
596
assert mock_bucket_not_exist .called
597
+ assert mock_get_project_number .called
600
598
assert (
601
599
output == "gs://test-project-vertex-pipelines-us-central1/output_artifacts/"
602
600
)
0 commit comments