diff --git a/olive/cli/shared_cache.py b/olive/cli/shared_cache.py index 78c89b5a5..e78d60507 100644 --- a/olive/cli/shared_cache.py +++ b/olive/cli/shared_cache.py @@ -32,12 +32,14 @@ def register_subcommand(parser): help="Confirm the deletion without prompting for confirmation.", ) sub_parser.add_argument( + "--account_name", "--account", type=str, required=True, help="The account name for the shared cache.", ) sub_parser.add_argument( + "--container_name", "--container", type=str, required=True, @@ -53,7 +55,7 @@ def register_subcommand(parser): @action def run(self): - container_client_factory = AzureContainerClientFactory(self.args.account, self.args.container) + container_client_factory = AzureContainerClientFactory(self.args.account_name, self.args.container_name) if self.args.delete: if self.args.all: if self.args.yes: @@ -63,4 +65,6 @@ def run(self): if confirm.lower() == "y": container_client_factory.delete_all() else: + if not self.args.model_hash: + raise ValueError("--model_hash is required when --delete is used without --all") container_client_factory.delete_blob(self.args.model_hash) diff --git a/test/cli/test_shared_cache.py b/test/cli/test_shared_cache.py new file mode 100644 index 000000000..00c6e6a89 --- /dev/null +++ b/test/cli/test_shared_cache.py @@ -0,0 +1,45 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +from argparse import ArgumentParser + +import pytest + +from olive.cli.shared_cache import SharedCacheCommand + + +def _parse_shared_cache_args(*args: str): + parser = ArgumentParser() + subparsers = parser.add_subparsers() + SharedCacheCommand.register_subcommand(subparsers) + return parser.parse_args(["shared-cache", *args]) + + +@pytest.mark.parametrize( + ("account_option", "container_option"), + [ + ("--account_name", "--container_name"), + ("--account", "--container"), + ], +) +def test_shared_cache_accepts_consistent_and_legacy_option_names(account_option: str, container_option: str): + args = _parse_shared_cache_args(account_option, "account", container_option, "container") + + assert args.account_name == "account" + assert args.container_name == "container" + + +def test_shared_cache_delete_requires_model_hash(): + args = _parse_shared_cache_args( + "--delete", + "--account_name", + "account", + "--container_name", + "container", + ) + command = object.__new__(SharedCacheCommand) + command.args = args + + with pytest.raises(ValueError, match="--model_hash is required"): + command.run()