fix for extract zip from gcs (kserve#3510)

* fix for extract zip from gcs

Signed-off-by: Andrews Arokiam <andrews.arokiam@ideas2it.com>

* initial commit for gcs model download unittests

Signed-off-by: Andrews Arokiam <andrews.arokiam@ideas2it.com>

* unittests for model download from gcs

Signed-off-by: Andrews Arokiam <andrews.arokiam@ideas2it.com>

* black format fix

Signed-off-by: Andrews Arokiam <andrews.arokiam@ideas2it.com>

* code verification

Signed-off-by: Andrews Arokiam <andrews.arokiam@ideas2it.com>

---------

Signed-off-by: Andrews Arokiam <andrews.arokiam@ideas2it.com>
Signed-off-by: asd981256 <asd981256@gmail.com>

python/kserve/kserve/storage/storage.py

@@ -302,7 +302,7 @@ def _download_gcs(uri, temp_dir: str):
                 dest_path = os.path.join(temp_dir, subdir_object_key)
                 logging.info("Downloading: %s", dest_path)
                 blob.download_to_filename(dest_path)
-            file_count += 1
+                file_count += 1
         if file_count == 0:
             raise RuntimeError("Failed to fetch model. No model found in %s." % uri)
 

python/kserve/kserve/storage/test/test_gcs_storage.py

@@ -0,0 +1,117 @@
+# Copyright 2024 The KServe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest.mock as mock
+import pytest
+from kserve.storage import Storage
+
+STORAGE_MODULE = "kserve.storage.storage"
+
+
+def get_call_args(call_args_list):
+    arg_list = []
+    for call in call_args_list:
+        args, _ = call
+        arg_list.append(args)
+    return arg_list
+
+
+def create_mock_dir(name):
+    mock_dir = mock.MagicMock()
+    mock_dir.name = name
+    return mock_dir
+
+
+def create_mock_dir_with_file(dir_name, file_name):
+    mock_obj = mock.MagicMock()
+    mock_obj.name = f"{dir_name}/{file_name}"
+    return mock_obj
+
+
+@mock.patch(STORAGE_MODULE + ".storage")
+def test_gcs_with_empty_dir(mock_storage):
+    gcs_path = "gs://foo/bar"
+
+    mock_storage.Client().bucket().list_blobs().__iter__.return_value = [
+        create_mock_dir("bar/")
+    ]
+
+    with pytest.raises(Exception):
+        Storage.download(gcs_path)
+
+
+@mock.patch(STORAGE_MODULE + ".storage")
+def test_gcs_with_nested_sub_dir(mock_storage):
+    gcs_path = "gs://foo/bar/test"
+
+    mock_root_dir = create_mock_dir("bar/")
+    mock_sub_dir = create_mock_dir("test/")
+    mock_file = create_mock_dir_with_file("test", "mock.object")
+
+    mock_storage.Client().bucket().list_blobs().__iter__.return_value = [
+        mock_root_dir,
+        mock_sub_dir,
+        mock_file,
+    ]
+    Storage.download(gcs_path)
+
+    arg_list = get_call_args(mock_file.download_to_filename.call_args_list)
+    assert "test/mock.object" in arg_list[0][0]
+
+
+@mock.patch(STORAGE_MODULE + ".storage")
+def test_download_model_from_gcs(mock_storage):
+    gcs_path = "gs://foo/bar"
+
+    mock_dir = create_mock_dir("bar/")
+    mock_file = create_mock_dir_with_file("bar", "mock.object")
+
+    mock_storage.Client().bucket().list_blobs().__iter__.return_value = [
+        mock_dir,
+        mock_file,
+    ]
+    Storage.download(gcs_path)
+
+    arg_list = get_call_args(mock_file.download_to_filename.call_args_list)
+    assert "/mock.object" in arg_list[0][0]
+
+
+@mock.patch("os.remove")
+@mock.patch("os.mkdir")
+@mock.patch("zipfile.ZipFile")
+@mock.patch(STORAGE_MODULE + ".storage")
+def test_gcs_model_unpack_archive_file(
+    mock_storage, MockZipFile, mock_create, mock_remove
+):
+    gcs_path = "gs://foo/bar"
+    output_dir = "test/out_dir"
+
+    mock_dir = create_mock_dir("bar/")
+    mock_file = create_mock_dir_with_file("bar", "mock.zip")
+    MockZipFile.return_value = mock_file
+
+    mock_storage.Client().bucket().list_blobs().__iter__.return_value = [
+        mock_dir,
+        mock_file,
+    ]
+    Storage.download(gcs_path, output_dir)
+
+    download_arg_list = get_call_args(mock_file.download_to_filename.call_args_list)
+
+    extract_arg_list = get_call_args(mock_file.extractall.call_args_list)
+
+    assert "/mock.zip" in download_arg_list[0][0]
+    assert output_dir == extract_arg_list[0][0]
+    assert mock_file.close.called
+    assert mock_remove.called

python/kserve/kserve/storage/test/test_storage.py

@@ -227,8 +227,13 @@ def test(_):
 def test_mock_gcs(mock_storage):
     gcs_path = "gs://foo/bar"
     mock_obj = mock.MagicMock()
-    mock_obj.name = "mock.object"
-    mock_storage.Client().bucket().list_blobs().__iter__.return_value = [mock_obj]
+    mock_obj.name = "bar/"
+    mock_obj1 = mock.MagicMock()
+    mock_obj1.name = "bar/mock.object"
+    mock_storage.Client().bucket().list_blobs().__iter__.return_value = [
+        mock_obj,
+        mock_obj1,
+    ]
     assert Storage.download(gcs_path)