Add suggestion from review and remove extra tests not relevant to this update

Author: holtskinner

google/cloud/documentai_toolbox/utilities/gcs_utilities.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 #
 """Google Cloud Storage utilities."""
-import importlib.metadata
 import os
 import re
 from typing import Dict, List, Optional, Tuple
@@ -143,18 +142,11 @@ def get_blob(
     if not re.match(constants.FILE_CHECK_REGEX, gcs_uri):
         raise ValueError("gcs_uri must link to a single file.")
 
-    try:
-        version = importlib.metadata.version("google-cloud-storage")
-    except importlib.metadata.PackageNotFoundError:
-        raise ImportError("google-cloud-storage is not installed.")
+    # google-cloud-storage >= 3.0.0
+    if hasattr(storage.Blob, "from_uri"):
+        return storage.Blob.from_uri(gcs_uri, _get_storage_client(module=module))
 
-    client = _get_storage_client(module=module)
-
-    major, _, _ = map(int, version.split("."))
-    if major < 3:
-        return storage.Blob.from_string(gcs_uri, client)
-    else:
-        return storage.Blob.from_uri(gcs_uri, client)
+    return storage.Blob.from_string(gcs_uri, _get_storage_client(module=module))
 
 
 def split_gcs_uri(gcs_uri: str) -> Tuple[str, str]:

tests/unit/test_gcs_utilities.py

@@ -598,72 +598,32 @@ def test_get_blobs_with_file_type_error():
         gcs_utilities.get_blobs(gcs_bucket_name="test-bucket", gcs_prefix="test.json")
 
 
-@mock.patch("google.cloud.documentai_toolbox.utilities.gcs_utilities.storage")
-def test_get_blob_success_major_3(mock_storage):
-    mock_version = "3.0.0"
-    with mock.patch("importlib.metadata.version", return_value=mock_version):
-        client = mock_storage.Client.return_value
-        gcs_uri = "gs://test-bucket/test.json"
-
-        gcs_utilities.get_blob(gcs_uri)
-
-        mock_storage.Blob.from_uri.assert_called_once_with(gcs_uri, client)
+def test_get_blob_invalid_uri():
+    with pytest.raises(ValueError, match="gcs_uri must link to a single file."):
+        gcs_utilities.get_blob("gs://test-bucket/prefix/")
 
 
 @mock.patch("google.cloud.documentai_toolbox.utilities.gcs_utilities.storage")
-def test_get_blob_success_major_2(mock_storage):
-    mock_version = "2.0.0"
-    with mock.patch("importlib.metadata.version", return_value=mock_version):
-        client = mock_storage.Client.return_value
-        gcs_uri = "gs://test-bucket/test.json"
-
-        gcs_utilities.get_blob(gcs_uri)
+def test_get_blob_from_uri(mock_storage):
+    gcs_uri = "gs://test-bucket/test.json"
 
-        mock_storage.Blob.from_string.assert_called_once_with(gcs_uri, client)
+    # Mock storage.Blob.from_uri to exist
+    mock_storage.Blob.from_uri.return_value = mock.Mock(spec=storage.blob.Blob)
 
+    gcs_utilities.get_blob(gcs_uri=gcs_uri)
 
-def test_get_blob_invalid_uri():
-    with pytest.raises(ValueError, match="gcs_uri must link to a single file."):
-        gcs_utilities.get_blob("gs://test-bucket/prefix/")
-
-
-def test_get_blob_import_error():
-    with mock.patch(
-        "importlib.metadata.version",
-        side_effect=importlib.metadata.PackageNotFoundError,
-    ):
-        with pytest.raises(ImportError, match="google-cloud-storage is not installed."):
-            gcs_utilities.get_blob("gs://test-bucket/test.json")
+    mock_storage.Blob.from_uri.assert_called_once()
 
 
 @mock.patch("google.cloud.documentai_toolbox.utilities.gcs_utilities.storage")
-def test_print_gcs_document_tree_with_skipping_files(mock_storage, capfd):
-    client = mock_storage.Client.return_value
-    mock_bucket = mock.Mock()
-    client.Bucket.return_value = mock_bucket
+def test_get_blob_from_string(mock_storage):
+    gcs_uri = "gs://test-bucket/test.json"
 
-    blobs = [
-        storage.Blob(
-            name=f"gs://test-directory/1/test_shard{i}.json",
-            bucket="gs://test-directory/1",
-        )
-        for i in range(1, 11)
-    ]
+    # Mock storage.Blob to NOT have from_uri
+    del mock_storage.Blob.from_uri
 
-    client.list_blobs.return_value = blobs
+    mock_storage.Blob.from_string.return_value = mock.Mock(spec=storage.blob.Blob)
 
-    # files_to_display = 2. 10 files total.
-    # idx 0, 1, 2 -> print
-    # idx 3, 4, 5, 6, 7, 8 -> skip
-    # idx 9 -> print last
-    gcs_utilities.print_gcs_document_tree(
-        gcs_bucket_name="test-directory", gcs_prefix="/", files_to_display=2
-    )
+    gcs_utilities.get_blob(gcs_uri=gcs_uri)
 
-    out, err = capfd.readouterr()
-    assert "├──test_shard1.json" in out
-    assert "├──test_shard2.json" in out
-    assert "├──test_shard3.json" in out
-    assert "├──test_shard4.json" not in out
-    assert "│  ...." in out
-    assert "└──test_shard10.json" in out
+    mock_storage.Blob.from_string.assert_called_once()