kserve/python/storage/test/test_azure_storage.py

# Copyright 2021 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
import shutil

from kserve_storage import Storage

STORAGE_MODULE = "kserve_storage.kserve_storage"


def create_mock_item(path):
    mock_obj = mock.MagicMock()
    mock_obj.name = path
    mock_obj.size = 1024  # Mock non-zero size for regular files

    # For async streaming, we need to mock the chunks() method
    async def mock_chunks():
        yield b"test"

    mock_stream = mock.MagicMock()
    mock_stream.chunks = mock_chunks
    mock_obj.readall.return_value = b"test"
    mock_obj.chunks = mock_chunks
    return mock_obj


def create_mock_blob(mock_storage, paths):
    mock_objs = [create_mock_item(path) for path in paths]
    mock_container = mock.MagicMock()

    # Mock async list_blobs (changed from walk_blobs)
    async def mock_list_blobs(name_starts_with=None, include=None):
        for obj in mock_objs:
            yield obj

    mock_container.list_blobs = mock_list_blobs

    # Mock async download_blob returning a stream (no blob_name param needed)
    async def mock_download_blob(max_concurrency=None):
        # Return a stream-like object with chunks() method
        stream = mock.MagicMock()

        async def mock_chunks():
            yield b"test"

        stream.chunks = mock_chunks
        return stream

    mock_blob_client = mock.MagicMock()
    mock_blob_client.download_blob = mock_download_blob
    mock_container.get_blob_client.return_value = mock_blob_client

    mock_svc = mock.MagicMock()
    mock_svc.get_container_client.return_value = mock_container

    # Make the service client work as an async context manager
    async def mock_aenter(self):
        return mock_svc

    async def mock_aexit(self, exc_type, exc_val, exc_tb):
        return None

    mock_svc.__aenter__ = mock_aenter
    mock_svc.__aexit__ = mock_aexit

    mock_storage.return_value = mock_svc
    return mock_storage, mock_container


def create_mock_dir(name):
    mock_file = mock.MagicMock()
    mock_file.name = name
    mock_file.is_directory = True
    return mock_file


def create_mock_file(name):
    mock_file = mock.MagicMock()
    mock_file.name = name
    mock_file.is_directory = False
    return mock_file


def create_mock_objects_for_file_share(mock_storage, mock_file_items):
    mock_share = mock.MagicMock()
    mock_share.list_directories_and_files.side_effect = mock_file_items

    mock_file = mock.MagicMock()
    mock_share.get_file_client.return_value = mock_file
    mock_data = mock.MagicMock()
    mock_file.download_file.return_value = mock_data

    mock_svc = mock.MagicMock()
    mock_svc.get_share_client.return_value = mock_share

    mock_storage.return_value = mock_svc
    return mock_storage, mock_share, mock_data


def get_call_args(call_args_list):
    arg_list = []
    for call in call_args_list:
        args, _ = call
        arg_list.append(args)
    return arg_list


# pylint: disable=protected-access
@pytest.fixture(scope="session", autouse=True)
def test_cleanup():
    yield None
    # Will be executed after the last test
    shutil.rmtree("some", ignore_errors=True)
    shutil.rmtree("dest_path", ignore_errors=True)


@mock.patch(STORAGE_MODULE + ".os.makedirs")
@mock.patch("azure.storage.blob.aio.BlobServiceClient")
def test_blob(mock_storage, mock_makedirs):  # pylint: disable=unused-argument
    # given
    blob_path = "https://kfserving.blob.core.windows.net/triton/simple_string/"
    paths = ["simple_string/1/model.graphdef", "simple_string/config.pbtxt"]
    mock_blob, mock_container = create_mock_blob(mock_storage, paths)

    # when
    Storage._download_azure_blob(blob_path, "dest_path")

    # then
    # Check that get_blob_client was called for each blob
    arg_list = get_call_args(mock_container.get_blob_client.call_args_list)
    assert set(arg_list) == set(
        [("simple_string/1/model.graphdef",), ("simple_string/config.pbtxt",)]
    )

    mock_storage.assert_called_with(
        "https://kfserving.blob.core.windows.net", credential=None
    )


@mock.patch(STORAGE_MODULE + ".os.makedirs")
@mock.patch("azure.storage.blob.aio.BlobServiceClient")
def test_blob_file_direct(mock_storage, mock_makedirs):  # pylint: disable=unused-argument
    # given
    blob_path = "https://accountname.blob.core.windows.net/container/somefile.text"
    paths = ["somefile.text"]
    mock_blob, mock_container = create_mock_blob(mock_storage, paths)

    # when
    Storage._download_azure_blob(blob_path, "dest_path")

    # then
    arg_list = get_call_args(mock_container.get_blob_client.call_args_list)
    assert arg_list == [("somefile.text",)]
    mock_storage.assert_called_with(
        "https://accountname.blob.core.windows.net", credential=None
    )


@mock.patch(STORAGE_MODULE + ".os.makedirs")
@mock.patch(STORAGE_MODULE + ".Storage._get_azure_storage_token")
@mock.patch("azure.storage.blob.aio.BlobServiceClient")
def test_secure_blob(mock_storage, mock_get_token, mock_makedirs):  # pylint: disable=unused-argument
    # given
    blob_path = "https://kfsecured.blob.core.windows.net/triton/simple_string/"
    mock_get_token.return_value = "some_token"
    # Set up empty blob mock to trigger RuntimeError
    mock_blob, mock_container = create_mock_blob(mock_storage, [])

    # when
    with pytest.raises(RuntimeError):
        Storage._download_azure_blob(blob_path, "dest_path")

    # then
    mock_get_token.assert_called()
    arg_list = []
    for call in mock_storage.call_args_list:
        _, kwargs = call
        arg_list.append(kwargs)
    assert arg_list == [{"credential": "some_token"}]


@mock.patch(STORAGE_MODULE + ".os.makedirs")
@mock.patch("azure.storage.blob.aio.BlobServiceClient")
def test_deep_blob(mock_storage, mock_makedirs):  # pylint: disable=unused-argument
    # given
    blob_path = (
        "https://accountname.blob.core.windows.net/container/some/deep/blob/path"
    )
    paths = ["f1", "f2", "d1/f11", "d1/d2/f21", "d1/d2/d3/f1231", "d4/f41"]
    fq_item_paths = ["some/deep/blob/path/" + p for p in paths]
    expected_calls = [(f,) for f in fq_item_paths]

    # when
    mock_blob, mock_container = create_mock_blob(mock_storage, fq_item_paths)
    try:
        Storage._download_azure_blob(blob_path, "some/dest/path")
    except OSError:  # Permissions Error Handling
        pass

    # then
    actual_calls = get_call_args(mock_container.get_blob_client.call_args_list)
    assert set(actual_calls) == set(expected_calls)


@mock.patch(STORAGE_MODULE + ".os.makedirs")
@mock.patch("azure.storage.blob.aio.BlobServiceClient")
def test_blob_file(mock_storage, mock_makedirs):  # pylint: disable=unused-argument
    # given
    blob_path = "https://accountname.blob.core.windows.net/container/somefile.text"
    paths = ["somefile"]
    fq_item_paths = paths
    expected_calls = [(f,) for f in fq_item_paths]

    # when
    mock_blob, mock_container = create_mock_blob(mock_storage, paths)
    Storage._download_azure_blob(blob_path, "some/dest/path")

    # then
    actual_calls = get_call_args(mock_container.get_blob_client.call_args_list)
    assert actual_calls == expected_calls


@mock.patch(STORAGE_MODULE + ".os.makedirs")
@mock.patch("azure.storage.blob.aio.BlobServiceClient")
def test_blob_fq_file(mock_storage, mock_makedirs):  # pylint: disable=unused-argument
    # given
    blob_path = (
        "https://accountname.blob.core.windows.net/container/folder/somefile.text"
    )
    paths = ["somefile"]
    fq_item_paths = ["folder/" + p for p in paths]
    expected_calls = [(f,) for f in fq_item_paths]

    # when
    mock_blob, mock_container = create_mock_blob(mock_storage, fq_item_paths)
    Storage._download_azure_blob(blob_path, "some/dest/path")

    # then
    actual_calls = get_call_args(mock_container.get_blob_client.call_args_list)
    assert actual_calls == expected_calls


@mock.patch(STORAGE_MODULE + ".os.makedirs")
@mock.patch("azure.storage.blob.aio.BlobServiceClient")
def test_blob_no_prefix(mock_storage, mock_makedirs):  # pylint: disable=unused-argument
    # given
    blob_path = "https://accountname.blob.core.windows.net/container/"
    paths = ["somefile.text", "somefolder/somefile.text"]
    fq_item_paths = ["" + p for p in paths]
    expected_calls = [(f,) for f in fq_item_paths]

    # when
    mock_blob, mock_container = create_mock_blob(mock_storage, fq_item_paths)
    Storage._download_azure_blob(blob_path, "some/dest/path")

    # then
    actual_calls = get_call_args(mock_container.get_blob_client.call_args_list)
    assert set(actual_calls) == set(expected_calls)


@mock.patch(STORAGE_MODULE + ".os.makedirs")
@mock.patch(STORAGE_MODULE + ".Storage._get_azure_storage_access_key")
@mock.patch("azure.storage.fileshare.ShareServiceClient")
def test_file_share(mock_storage, mock_get_access_key, mock_makedirs):  # pylint: disable=unused-argument
    # given
    file_share_path = "https://kfserving.file.core.windows.net/triton/simple_string/"
    mock_get_access_key.return_value = "some_token"

    mock_file_share, mock_file, mock_data = create_mock_objects_for_file_share(
        mock_storage,
        [
            [create_mock_dir("1"), create_mock_file("config.pbtxt")],
            [create_mock_file("model.graphdef")],
            [],
        ],
    )

    # when
    Storage._download_azure_file_share(file_share_path, "dest_path")

    # then
    arg_list = get_call_args(mock_file.get_file_client.call_args_list)
    assert set(arg_list) == set(
        [("simple_string/1/model.graphdef",), ("simple_string/config.pbtxt",)]
    )

    # then
    mock_get_access_key.assert_called()
    mock_storage.assert_called_with(
        "https://kfserving.file.core.windows.net", credential="some_token"
    )


@mock.patch(STORAGE_MODULE + ".os.makedirs")
@mock.patch(STORAGE_MODULE + ".Storage._get_azure_storage_access_key")
@mock.patch("azure.storage.fileshare.ShareServiceClient")
def test_deep_file_share(mock_storage, mock_get_access_key, mock_makedirs):  # pylint: disable=unused-argument
    file_share_path = (
        "https://accountname.file.core.windows.net/container/some/deep/blob/path"
    )
    paths = ["f1", "f2", "d1/f11", "d1/d2/f21", "d1/d2/d3/f1231", "d4/f41"]
    fq_item_paths = ["some/deep/blob/path/" + p for p in paths]
    expected_calls = [(f,) for f in fq_item_paths]
    mock_get_access_key.return_value = "some_token"

    # when
    mock_file_share, mock_file, mock_data = create_mock_objects_for_file_share(
        mock_storage,
        [
            [
                create_mock_dir("d1"),
                create_mock_dir("d4"),
                create_mock_file("f1"),
                create_mock_file("f2"),
            ],
            [create_mock_file("f41")],
            [create_mock_dir("d2"), create_mock_file("f11")],
            [create_mock_dir("d3"), create_mock_file("f21")],
            [create_mock_file("f1231")],
            [],
        ],
    )
    try:
        Storage._download_azure_file_share(file_share_path, "some/dest/path")
    except OSError:  # Permissions Error Handling
        pass

    # then
    actual_calls = get_call_args(mock_file.get_file_client.call_args_list)
    assert set(actual_calls) == set(expected_calls)

    # then
    mock_get_access_key.assert_called()
    mock_storage.assert_called_with(
        "https://accountname.file.core.windows.net", credential="some_token"
    )


@mock.patch(STORAGE_MODULE + ".os.makedirs")
@mock.patch(STORAGE_MODULE + ".Storage._get_azure_storage_access_key")
@mock.patch("azure.storage.fileshare.ShareServiceClient")
def test_file_share_fq_file(mock_storage, mock_get_access_key, mock_makedirs):  # pylint: disable=unused-argument
    # given
    file_share_path = "https://accountname.file.core.windows.net/container/folder/"
    paths = ["somefile.text"]
    fq_item_paths = ["folder/" + p for p in paths]
    expected_calls = [(f,) for f in fq_item_paths]

    # when
    mock_get_access_key.return_value = "some_token"
    mock_file_share, mock_file, mock_data = create_mock_objects_for_file_share(
        mock_storage, [[create_mock_file("somefile.text")], []]
    )
    Storage._download_azure_file_share(file_share_path, "some/dest/path")

    # then
    actual_calls = get_call_args(mock_file.get_file_client.call_args_list)
    assert actual_calls == expected_calls

    # then
    mock_get_access_key.assert_called()
    mock_storage.assert_called_with(
        "https://accountname.file.core.windows.net", credential="some_token"
    )


@mock.patch(STORAGE_MODULE + ".os.makedirs")
@mock.patch(STORAGE_MODULE + ".Storage._get_azure_storage_access_key")
@mock.patch("azure.storage.fileshare.ShareServiceClient")
def test_file_share_no_prefix(mock_storage, mock_get_access_key, mock_makedirs):  # pylint: disable=unused-argument
    # given
    file_share_path = "https://accountname.file.core.windows.net/container/"
    paths = ["somefile.text", "somefolder/somefile.text"]
    fq_item_paths = ["" + p for p in paths]
    expected_calls = [(f,) for f in fq_item_paths]

    # when
    mock_get_access_key.return_value = "some_token"
    mock_file_share, mock_file, mock_data = create_mock_objects_for_file_share(
        mock_storage,
        [
            [create_mock_dir("somefolder"), create_mock_file("somefile.text")],
            [create_mock_file("somefile.text")],
            [],
        ],
    )
    Storage._download_azure_file_share(file_share_path, "some/dest/path")

    # then
    arg_list = get_call_args(mock_file.get_file_client.call_args_list)
    assert set(arg_list) == set(expected_calls)

    # then
    mock_get_access_key.assert_called()
    mock_storage.assert_called_with(
        "https://accountname.file.core.windows.net", credential="some_token"
    )


@mock.patch(STORAGE_MODULE + ".os.makedirs")
@mock.patch("azure.storage.blob.aio.BlobServiceClient")
def test_blob_allow_patterns(mock_storage, mock_makedirs):
    # given
    blob_path = "https://kfserving.blob.core.windows.net/triton/model/"
    paths = [
        "model/weights.safetensors",
        "model/weights.bin",
        "model/config.json",
    ]
    mock_blob, mock_container = create_mock_blob(mock_storage, paths)

    # when
    Storage._download_azure_blob(
        blob_path, "dest_path", allow_patterns=["*.safetensors", "*.json"]
    )

    # then
    arg_list = get_call_args(mock_container.get_blob_client.call_args_list)
    blob_names = [a[0] for a in arg_list]
    assert "model/weights.safetensors" in blob_names
    assert "model/config.json" in blob_names
    assert "model/weights.bin" not in blob_names


@mock.patch(STORAGE_MODULE + ".os.makedirs")
@mock.patch("azure.storage.blob.aio.BlobServiceClient")
def test_blob_ignore_patterns(mock_storage, mock_makedirs):
    # given
    blob_path = "https://kfserving.blob.core.windows.net/triton/model/"
    paths = [
        "model/weights.safetensors",
        "model/weights.bin",
        "model/config.json",
    ]
    mock_blob, mock_container = create_mock_blob(mock_storage, paths)

    # when
    Storage._download_azure_blob(blob_path, "dest_path", ignore_patterns=["*.bin"])

    # then
    arg_list = get_call_args(mock_container.get_blob_client.call_args_list)
    blob_names = [a[0] for a in arg_list]
    assert "model/weights.bin" not in blob_names
    assert len(blob_names) == 2


@mock.patch(STORAGE_MODULE + ".os.makedirs")
@mock.patch(STORAGE_MODULE + ".Storage._get_azure_storage_access_key")
@mock.patch("azure.storage.fileshare.ShareServiceClient")
def test_file_share_allow_patterns(mock_storage, mock_get_access_key, mock_makedirs):
    # given
    file_share_path = "https://kfserving.file.core.windows.net/triton/model/"
    mock_get_access_key.return_value = "some_token"

    mock_file_share, mock_file, mock_data = create_mock_objects_for_file_share(
        mock_storage,
        [
            [
                create_mock_file("weights.safetensors"),
                create_mock_file("weights.bin"),
                create_mock_file("config.json"),
            ],
            [],
        ],
    )

    # when
    Storage._download_azure_file_share(
        file_share_path, "dest_path", allow_patterns=["*.safetensors", "*.json"]
    )

    # then
    arg_list = get_call_args(mock_file.get_file_client.call_args_list)
    file_paths = [a[0] for a in arg_list]
    assert "model/weights.safetensors" in file_paths
    assert "model/config.json" in file_paths
    assert "model/weights.bin" not in file_paths


@mock.patch(STORAGE_MODULE + ".os.makedirs")
@mock.patch(STORAGE_MODULE + ".Storage._get_azure_storage_access_key")
@mock.patch("azure.storage.fileshare.ShareServiceClient")
def test_file_share_ignore_patterns(mock_storage, mock_get_access_key, mock_makedirs):
    # given
    file_share_path = "https://kfserving.file.core.windows.net/triton/model/"
    mock_get_access_key.return_value = "some_token"

    mock_file_share, mock_file, mock_data = create_mock_objects_for_file_share(
        mock_storage,
        [
            [
                create_mock_file("weights.safetensors"),
                create_mock_file("weights.bin"),
                create_mock_file("config.json"),
            ],
            [],
        ],
    )

    # when
    Storage._download_azure_file_share(
        file_share_path, "dest_path", ignore_patterns=["*.bin"]
    )

    # then
    arg_list = get_call_args(mock_file.get_file_client.call_args_list)
    file_paths = [a[0] for a in arg_list]
    assert "model/weights.bin" not in file_paths
    assert len(file_paths) == 2