@@ -579,3 +579,50 @@ def test_get_blobs_with_no_input():
579579 match = "You must provide either `gcs_uri` or both `gcs_bucket_name` and `gcs_prefix`." ,
580580 ):
581581 gcs_utilities .get_blobs ()
582+
583+
584+ @mock .patch ("google.cloud.documentai_toolbox.utilities.gcs_utilities.storage" )
585+ def test_get_blobs_with_gcs_uri (mock_storage ):
586+ client = mock_storage .Client .return_value
587+ gcs_uri = "gs://test-bucket/test-directory/1/"
588+
589+ gcs_utilities .get_blobs (gcs_uri = gcs_uri )
590+
591+ mock_storage .Client .assert_called_once ()
592+ client .list_blobs .assert_called_once_with ("test-bucket" , prefix = "test-directory/1/" )
593+
594+
595+ def test_get_blobs_with_file_type_error ():
596+ with pytest .raises (ValueError , match = "gcs_prefix cannot contain file types" ):
597+ gcs_utilities .get_blobs (gcs_bucket_name = "test-bucket" , gcs_prefix = "test.json" )
598+
599+
600+ def test_get_blob_invalid_uri ():
601+ with pytest .raises (ValueError , match = "gcs_uri must link to a single file." ):
602+ gcs_utilities .get_blob ("gs://test-bucket/prefix/" )
603+
604+
605+ @mock .patch ("google.cloud.documentai_toolbox.utilities.gcs_utilities.storage" )
606+ def test_get_blob_from_uri (mock_storage ):
607+ gcs_uri = "gs://test-bucket/test.json"
608+
609+ # Mock storage.Blob.from_uri to exist
610+ mock_storage .Blob .from_uri .return_value = mock .Mock (spec = storage .blob .Blob )
611+
612+ gcs_utilities .get_blob (gcs_uri = gcs_uri )
613+
614+ mock_storage .Blob .from_uri .assert_called_once ()
615+
616+
617+ @mock .patch ("google.cloud.documentai_toolbox.utilities.gcs_utilities.storage" )
618+ def test_get_blob_from_string (mock_storage ):
619+ gcs_uri = "gs://test-bucket/test.json"
620+
621+ # Mock storage.Blob to NOT have from_uri
622+ del mock_storage .Blob .from_uri
623+
624+ mock_storage .Blob .from_string .return_value = mock .Mock (spec = storage .blob .Blob )
625+
626+ gcs_utilities .get_blob (gcs_uri = gcs_uri )
627+
628+ mock_storage .Blob .from_string .assert_called_once ()
0 commit comments