diff --git a/cfa/cloudops/scripts.py b/cfa/cloudops/scripts.py index 6a20bd3..e132612 100644 --- a/cfa/cloudops/scripts.py +++ b/cfa/cloudops/scripts.py @@ -1176,7 +1176,7 @@ def add_tasks_from_yaml(): help="Base command for the tasks", ) parser.add_argument( - "-f", + "-fp", "--file_path", type=str, required=True, diff --git a/changelog.md b/changelog.md index 97560a7..b0aece0 100644 --- a/changelog.md +++ b/changelog.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/). The versioning pattern is `major.minor.patch`. --- +## 0.7.8 + +- added tests for increased coverage + ## 0.7.7 - added `create_new_folder` flag for upload functions to allow/warn users when uploading to a folder that does not currently exist in Blob. diff --git a/pyproject.toml b/pyproject.toml index c3e32d0..f89129e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "cfa.cloudops" -version = "0.7.7" +version = "0.7.8" description = "Cloud storage, batch, functions, MLOps assistance" authors = [ {name = "Ryan Raasch", email = "xng3@cdc.gov"} diff --git a/tests/test_auth_and_util_more.py b/tests/test_auth_and_util_more.py new file mode 100644 index 0000000..efd9ad9 --- /dev/null +++ b/tests/test_auth_and_util_more.py @@ -0,0 +1,488 @@ +import os +from types import SimpleNamespace + +import pytest + +from cfa.cloudops import auth, util + + +def test_lookup_service_principal_success(monkeypatch): + payload = '[{"appId": "abc"}]' + + monkeypatch.setattr("cfa.cloudops.util.sp.check_output", lambda *a, **k: payload) + + result = util.lookup_service_principal("my-sp") + assert result == [{"appId": "abc"}] + + +def test_lookup_service_principal_failure(monkeypatch): + def boom(*args, **kwargs): + raise RuntimeError("az failed") + + monkeypatch.setattr("cfa.cloudops.util.sp.check_output", boom) + + with pytest.raises(RuntimeError): + util.lookup_service_principal("my-sp") + + +def test_lookup_available_vm_skus_for_batch_to_dict(monkeypatch): + sku1 = SimpleNamespace( + name="Standard_D2s_v3", + family_name="fam", + batch_support_end_of_life=None, + additional_properties={"tier": "standard"}, + capabilities=[SimpleNamespace(name="vCPUs", value="2")], + ) + sku2 = SimpleNamespace( + name="Standard_D4s_v3", + family_name="fam", + batch_support_end_of_life=None, + additional_properties={}, + capabilities=[], + ) + + client = SimpleNamespace( + location=SimpleNamespace( + list_supported_virtual_machine_skus=lambda **kwargs: [sku1, sku2] + ) + ) + + monkeypatch.setattr("cfa.cloudops.util.get_config_val", lambda *a, **k: "eastus") + + result = util.lookup_available_vm_skus_for_batch(client=client, to_dict=True) + + assert len(result) == 2 + assert result[0]["name"] == "Standard_D2s_v3" + assert result[0]["vCPUs"] == "2" + + +def test_lookup_available_vm_skus_for_batch_builds_client(monkeypatch): + sku = SimpleNamespace( + name="Standard_D2s_v3", + family_name="fam", + batch_support_end_of_life=None, + additional_properties={}, + capabilities=[], + ) + + client = SimpleNamespace( + location=SimpleNamespace( + list_supported_virtual_machine_skus=lambda **kwargs: [sku] + ) + ) + + monkeypatch.setattr( + "cfa.cloudops.client.get_batch_management_client", + lambda **kwargs: client, + ) + monkeypatch.setattr("cfa.cloudops.util.get_config_val", lambda *a, **k: "eastus") + + result = util.lookup_available_vm_skus_for_batch(client=None, to_dict=False) + assert result == [sku] + + +def test_credential_handler_require_attr(): + ch = auth.CredentialHandler() + + with pytest.raises(AttributeError) as exc: + ch.require_attr(["azure_tenant_id", "azure_client_id"], goal="auth") + + assert "azure_tenant_id" in str(exc.value) + assert "azure_client_id" in str(exc.value) + + +def test_credential_handler_endpoint_properties(): + ch = auth.CredentialHandler( + azure_batch_account="acct", + azure_batch_location="eastus", + azure_batch_endpoint_subdomain="batch.azure.com/", + azure_blob_storage_account="blobacct", + azure_blob_storage_endpoint_subdomain="blob.core.windows.net/", + azure_container_registry_account="reg", + azure_container_registry_domain="azurecr.io", + ) + + assert ch.azure_batch_endpoint == "https://acct.eastus.batch.azure.com/" + assert ch.azure_blob_storage_endpoint == "https://blobacct.blob.core.windows.net/" + assert ch.azure_container_registry_endpoint == "https://reg.azurecr.io" + + +def test_credential_handler_user_credential(monkeypatch): + sentinel = object() + monkeypatch.setattr("cfa.cloudops.auth.ManagedIdentityCredential", lambda: sentinel) + + ch = auth.CredentialHandler() + assert ch.user_credential is sentinel + + +def test_service_principal_secret_branches(monkeypatch): + monkeypatch.setattr("cfa.cloudops.auth.get_sp_secret", lambda *a, **k: "kv-secret") + + ch_sp = auth.CredentialHandler( + azure_keyvault_endpoint="https://kv", + azure_keyvault_sp_secret_id="sp-id", + method="sp", + ) + ch_sp.azure_client_secret = "direct-secret" # pragma: allowlist secret + assert ch_sp.service_principal_secret == "direct-secret" # pragma: allowlist secret + + ch_default = auth.CredentialHandler( + azure_keyvault_endpoint="https://kv", + azure_keyvault_sp_secret_id="sp-id", + method="default", + ) + ch_default.__dict__["default_credential"] = "default-cred" + assert ( + ch_default.service_principal_secret == "kv-secret" # pragma: allowlist secret + ) + + ch_env = auth.CredentialHandler( + azure_keyvault_endpoint="https://kv", + azure_keyvault_sp_secret_id="sp-id", + method="env", + ) + ch_env.__dict__["user_credential"] = "user-cred" + assert ch_env.service_principal_secret == "kv-secret" # pragma: allowlist secret + + +def test_batch_service_principal_credentials(monkeypatch): + called = {} + + def fake_spcred(**kwargs): + called.update(kwargs) + return SimpleNamespace(**kwargs) + + monkeypatch.setattr("cfa.cloudops.auth.ServicePrincipalCredentials", fake_spcred) + + ch = auth.CredentialHandler( + azure_tenant_id="tenant", + azure_client_id="client", + azure_batch_resource_url="resource", + ) + ch.__dict__["service_principal_secret"] = "secret" # pragma: allowlist secret + + cred = ch.batch_service_principal_credentials + assert cred.client_id == "client" + assert called["secret"] == "secret" # pragma: allowlist secret + + +def test_client_secret_credential_variants(monkeypatch): + calls = [] + + def fake_client_secret_cred(**kwargs): + calls.append(kwargs) + return SimpleNamespace(**kwargs) + + monkeypatch.setattr( + "cfa.cloudops.auth.ClientSecretCredential", fake_client_secret_cred + ) + + ch = auth.CredentialHandler(azure_tenant_id="t", azure_client_id="c") + ch.__dict__["service_principal_secret"] = "s1" # pragma: allowlist secret + out1 = ch.client_secret_sp_credential + assert out1.client_secret == "s1" # pragma: allowlist secret + + ch2 = auth.CredentialHandler(azure_tenant_id="t", azure_client_id="c") + ch2.azure_client_secret = "s2" # pragma: allowlist secret + out2 = ch2.client_secret_credential + assert out2.client_secret == "s2" # pragma: allowlist secret + + +def test_compute_node_identity_reference(): + ch = auth.CredentialHandler( + azure_user_assigned_identity="/subscriptions/sub/resourceGroups/rg/providers/id" + ) + ref = ch.compute_node_identity_reference + assert ref.resource_id.endswith("/providers/id") + + +def test_azure_container_registry_valid_and_invalid(monkeypatch): + ch = auth.CredentialHandler( + azure_container_registry_account="reg", + azure_container_registry_domain="azurecr.io", + azure_user_assigned_identity="/subscriptions/sub/resourceGroups/rg/providers/id", + ) + + monkeypatch.setattr( + "cfa.cloudops.auth.is_valid_acr_endpoint", lambda endpoint: (True, None) + ) + monkeypatch.setattr( + "cfa.cloudops.auth.batch_mgmt_models.ContainerRegistry", + lambda **kwargs: SimpleNamespace(**kwargs), + ) + + reg = ch.azure_container_registry + assert reg.user_name == "reg" + + ch2 = auth.CredentialHandler( + azure_container_registry_account="reg", + azure_container_registry_domain="azurecr.io", + azure_user_assigned_identity="/subscriptions/sub/resourceGroups/rg/providers/id", + ) + monkeypatch.setattr( + "cfa.cloudops.auth.is_valid_acr_endpoint", + lambda endpoint: (False, "bad endpoint"), + ) + with pytest.raises(ValueError): + _ = ch2.azure_container_registry + + +def test_default_credential_wrapper(monkeypatch): + class FakeCredential: + def get_token(self, *scopes, **kwargs): + return "tok" + + class FakePolicy: + def __init__(self, credential, resource_id, **kwargs): + self.credential = credential + self.resource_id = resource_id + + def on_request(self, request): + request.http_request.headers["Authorization"] = "Bearer abc123" + + monkeypatch.setattr("cfa.cloudops.auth.BearerTokenCredentialPolicy", FakePolicy) + + dc = auth.DefaultCredential(credential=FakeCredential()) + assert dc.get_token("scope") == "tok" + dc.set_token() + assert dc.token["access_token"] == "abc123" + + +def test_get_sp_secret(monkeypatch): + monkeypatch.setattr( + "cfa.cloudops.auth.ManagedIdentityCredential", lambda: "managed" + ) + monkeypatch.setattr( + "cfa.cloudops.auth.SecretClient", + lambda vault_url, credential: SimpleNamespace( + get_secret=lambda sid: SimpleNamespace(value=f"secret-{sid}") + ), + ) + + result = auth.get_sp_secret("https://kv", "sp-id") + assert result == "secret-sp-id" + + +def test_get_client_secret_sp_credential(monkeypatch): + monkeypatch.setattr("cfa.cloudops.auth.get_sp_secret", lambda *a, **k: "sp-secret") + monkeypatch.setattr( + "cfa.cloudops.auth.ClientSecretCredential", + lambda **kwargs: SimpleNamespace(**kwargs), + ) + + cred = auth.get_client_secret_sp_credential( + vault_url="https://kv", + vault_sp_secret_id="sp-id", + tenant_id="tenant", + application_id="app", + ) + assert cred.client_secret == "sp-secret" # pragma: allowlist secret + + +def test_get_service_principal_credentials(monkeypatch): + monkeypatch.setattr("cfa.cloudops.auth.get_sp_secret", lambda *a, **k: "sp-secret") + monkeypatch.setattr( + "cfa.cloudops.auth.ServicePrincipalCredentials", + lambda **kwargs: SimpleNamespace(**kwargs), + ) + + cred = auth.get_service_principal_credentials( + vault_url="https://kv", + vault_sp_secret_id="sp-id", + tenant_id="tenant", + application_id="app", + ) + assert cred.secret == "sp-secret" # pragma: allowlist secret + + +def test_get_compute_node_identity_reference_helper(monkeypatch): + identity = SimpleNamespace(resource_id="rid") + + monkeypatch.setattr( + "cfa.cloudops.auth.EnvCredentialHandler", + lambda: SimpleNamespace(compute_node_identity_reference=identity), + ) + + result = auth.get_compute_node_identity_reference() + assert result is identity + + +def test_get_secret_client(monkeypatch): + captured = {} + + def fake_secret_client(vault_url, credential): + captured["vault_url"] = vault_url + captured["credential"] = credential + return SimpleNamespace(vault_url=vault_url) + + monkeypatch.setattr("cfa.cloudops.auth.SecretClient", fake_secret_client) + + client = auth.get_secret_client("mykv", credential="cred") + assert client.vault_url == "https://mykv.vault.azure.net" + assert captured["credential"] == "cred" + + +def test_load_keyvault_vars_force_and_skip(monkeypatch): + class FakeSecretClient: + def __init__(self): + self.calls = [] + + def get_secret(self, key): + self.calls.append(key) + return SimpleNamespace(value=f"value-{key}") + + sc = FakeSecretClient() + + monkeypatch.setenv("AZURE_BATCH_ACCOUNT", "existing") + + auth.load_keyvault_vars(sc, force_keyvault=False) + assert os.environ["AZURE_BATCH_ACCOUNT"] == "existing" + + auth.load_keyvault_vars(sc, force_keyvault=True) + assert "AZURE-BATCH-ACCOUNT" in [c.upper() for c in sc.calls] + + +def test_get_keyvault_vars_none_and_success(monkeypatch): + assert auth.get_keyvault_vars(None, credential="cred") is None + + seen = {} + + monkeypatch.setattr( + "cfa.cloudops.auth.get_secret_client", + lambda keyvault, credential: "secret-client", + ) + + def fake_load(secret_client, force_keyvault=False): + seen["client"] = secret_client + seen["force"] = force_keyvault + + monkeypatch.setattr("cfa.cloudops.auth.load_keyvault_vars", fake_load) + + auth.get_keyvault_vars("mykv", credential="cred", force_keyvault=True) + assert seen["client"] == "secret-client" + assert seen["force"] is True + + +def test_load_env_vars(monkeypatch): + class FakeSub: + subscription_id = "sub-1" + tenant_id = "tenant-1" + display_name = "rg-name" + + monkeypatch.setattr("cfa.cloudops.auth.load_dotenv", lambda *a, **k: None) + monkeypatch.setattr("cfa.cloudops.auth.ManagedIdentityCredential", lambda: "mid") + monkeypatch.setattr( + "cfa.cloudops.auth.SubscriptionClient", + lambda cred: SimpleNamespace( + subscriptions=SimpleNamespace(list=lambda: [FakeSub()]) + ), + ) + + called = {"set_env": 0, "kv": 0} + monkeypatch.setattr( + "cfa.cloudops.auth.d.set_env_vars", + lambda: called.__setitem__("set_env", called["set_env"] + 1), + ) + monkeypatch.setattr( + "cfa.cloudops.auth.get_keyvault_vars", + lambda **kwargs: called.__setitem__("kv", called["kv"] + 1), + ) + + auth.load_env_vars( + dotenv_path=".env.test", keyvault_name="mykv", force_keyvault=True + ) + + assert os.environ["AZURE_SUBSCRIPTION_ID"] == "sub-1" + assert os.environ["AZURE_TENANT_ID"] == "tenant-1" + assert os.environ["AZURE_RESOURCE_GROUP_NAME"] == "rg-name" + assert called["set_env"] == 1 + assert called["kv"] == 1 + + +def test_env_credential_handler_init(monkeypatch): + monkeypatch.setattr("cfa.cloudops.auth.load_env_vars", lambda **kwargs: None) + monkeypatch.setattr( + "cfa.cloudops.auth.get_config_val", + lambda key, config_dict=None, try_env=True: ( + config_dict.get(key) + if config_dict and key in config_dict + else os.getenv(key.upper()) + ), + ) + + monkeypatch.delenv("AZURE_BATCH_LOCATION", raising=False) + + handler = auth.EnvCredentialHandler( + dotenv_path=".env.test", azure_batch_account="acct" + ) + assert handler.method == "env" + assert handler.azure_batch_location == auth.d.default_azure_batch_location + + +def test_sp_credential_handler_init(monkeypatch): + monkeypatch.setattr("cfa.cloudops.auth.load_dotenv", lambda *a, **k: None) + monkeypatch.setattr("cfa.cloudops.auth.d.set_env_vars", lambda: None) + monkeypatch.setattr("cfa.cloudops.auth.get_keyvault_vars", lambda **kwargs: None) + monkeypatch.setattr( + "cfa.cloudops.auth.ClientSecretCredential", + lambda **kwargs: SimpleNamespace(**kwargs), + ) + monkeypatch.setattr( + "cfa.cloudops.auth.get_config_val", + lambda key, config_dict=None, try_env=True: ( + config_dict.get(key) + if config_dict and key in config_dict + else os.getenv(key.upper()) + ), + ) + + handler = auth.SPCredentialHandler( + azure_tenant_id="tenant", + azure_subscription_id="sub", + azure_client_id="client", + azure_client_secret="secret", # pragma: allowlist secret + azure_batch_account="acct", + ) + assert handler.method == "sp" + + +def test_default_credential_handler_success(monkeypatch): + class FakeSub: + subscription_id = "sub-1" + display_name = "rg-name" + + monkeypatch.setenv("AZURE_SUBSCRIPTION_ID", "sub-1") + monkeypatch.setattr("cfa.cloudops.auth.load_dotenv", lambda *a, **k: None) + monkeypatch.setattr("cfa.cloudops.auth.d.set_env_vars", lambda: None) + monkeypatch.setattr("cfa.cloudops.auth.get_keyvault_vars", lambda **kwargs: None) + monkeypatch.setattr("cfa.cloudops.auth.DefaultCredential", lambda: "dcred") + monkeypatch.setattr( + "cfa.cloudops.auth.SubscriptionClient", + lambda cred: SimpleNamespace( + subscriptions=SimpleNamespace(list=lambda: [FakeSub()]) + ), + ) + monkeypatch.setattr( + "cfa.cloudops.auth.get_config_val", + lambda key, config_dict=None, try_env=True: ( + config_dict.get(key) + if config_dict and key in config_dict + else os.getenv(key.upper()) + ), + ) + + handler = auth.DefaultCredentialHandler(dotenv_path=".env.test") + assert handler.method == "default" + + +def test_default_credential_handler_missing_sub(monkeypatch): + monkeypatch.delenv("AZURE_SUBSCRIPTION_ID", raising=False) + monkeypatch.setattr("cfa.cloudops.auth.load_dotenv", lambda *a, **k: None) + monkeypatch.setattr("cfa.cloudops.auth.DefaultCredential", lambda: "dcred") + monkeypatch.setattr( + "cfa.cloudops.auth.SubscriptionClient", + lambda cred: SimpleNamespace(subscriptions=SimpleNamespace(list=lambda: [])), + ) + + with pytest.raises(ValueError): + auth.DefaultCredentialHandler(dotenv_path=".env.test") diff --git a/tests/test_automation.py b/tests/test_automation.py new file mode 100644 index 0000000..f4e407d --- /dev/null +++ b/tests/test_automation.py @@ -0,0 +1,349 @@ +from types import SimpleNamespace + +from cfa.cloudops import automation + + +class FakeClient: + def __init__(self, dotenv_path=None): + self.dotenv_path = dotenv_path + self.cred = SimpleNamespace( + azure_resource_group_name="rg", + azure_batch_account="acct", + ) + self.batch_mgmt_client = "bmc" + self.calls = { + "upload_folders": [], + "upload_files": [], + "create_job": [], + "add_task": [], + "add_tasks_from_yaml": [], + "monitor_job": [], + } + + def upload_folders(self, **kwargs): + self.calls["upload_folders"].append(kwargs) + + def upload_files(self, **kwargs): + self.calls["upload_files"].append(kwargs) + + def create_job(self, **kwargs): + self.calls["create_job"].append(kwargs) + + def add_task(self, **kwargs): + self.calls["add_task"].append(kwargs) + return f"tid-{len(self.calls['add_task'])}" + + def add_tasks_from_yaml(self, **kwargs): + self.calls["add_tasks_from_yaml"].append(kwargs) + + def monitor_job(self, job_name): + self.calls["monitor_job"].append(job_name) + + +def test_run_experiment_returns_none_when_client_creation_fails(monkeypatch): + monkeypatch.setattr( + automation, + "CloudClient", + lambda **kwargs: (_ for _ in ()).throw(RuntimeError("boom")), + ) + monkeypatch.setattr( + automation.toml, + "load", + lambda _: { + "job": {"pool_name": "pool", "job_name": "job"}, + "experiment": {"base_cmd": "echo hi", "x": [1]}, + }, + ) + + assert automation.run_experiment("exp.toml") is None + + +def test_run_experiment_returns_none_without_pool_name(monkeypatch): + fake = FakeClient() + monkeypatch.setattr(automation, "CloudClient", lambda **kwargs: fake) + monkeypatch.setattr( + automation.toml, "load", lambda _: {"job": {}, "experiment": {}} + ) + + assert automation.run_experiment("exp.toml") is None + + +def test_run_experiment_returns_none_when_pool_missing(monkeypatch): + fake = FakeClient() + monkeypatch.setattr(automation, "CloudClient", lambda **kwargs: fake) + monkeypatch.setattr( + automation.toml, + "load", + lambda _: { + "job": {"pool_name": "pool-x", "job_name": "job"}, + "experiment": {"base_cmd": "echo hi", "x": [1]}, + }, + ) + monkeypatch.setattr( + automation.batch_helpers, "check_pool_exists", lambda **kwargs: False + ) + + assert automation.run_experiment("exp.toml") is None + + +def test_run_experiment_yaml_upload_and_monitor(monkeypatch): + fake = FakeClient() + monkeypatch.setattr(automation, "CloudClient", lambda **kwargs: fake) + monkeypatch.setattr( + automation.batch_helpers, "check_pool_exists", lambda **kwargs: True + ) + monkeypatch.setattr( + automation.toml, + "load", + lambda _: { + "job": { + "pool_name": "pool-x", + "job_name": "job-1", + "save_logs_to_blob": "logs", + "logs_folder": "run1", + "task_retries": 2, + "container": "img:tag", + "monitor_job": True, + }, + "upload": { + "container_name": "data-cont", + "location_in_blob": "inputs", + "folders": ["src"], + "files": ["a.txt"], + }, + "experiment": { + "base_cmd": "python task.py", + "exp_yaml": "grid.yaml", + }, + }, + ) + + assert automation.run_experiment("exp.toml", dotenv_path=".env") is None + + assert fake.calls["upload_folders"][0]["folder_names"] == ["src"] + assert fake.calls["upload_folders"][0]["location_in_blob"] == "inputs" + assert fake.calls["upload_files"][0]["files"] == ["a.txt"] + assert fake.calls["create_job"][0]["job_name"] == "job-1" + assert fake.calls["create_job"][0]["task_retries"] == 2 + assert fake.calls["add_tasks_from_yaml"][0]["file_path"] == "grid.yaml" + assert fake.calls["monitor_job"] == ["job-1"] + + +def test_run_experiment_parameter_grid_adds_tasks(monkeypatch): + fake = FakeClient() + monkeypatch.setattr(automation, "CloudClient", lambda **kwargs: fake) + monkeypatch.setattr( + automation.batch_helpers, "check_pool_exists", lambda **kwargs: True + ) + monkeypatch.setattr( + automation.toml, + "load", + lambda _: { + "job": { + "pool_name": "pool-x", + "job_name": "job-2", + "monitor_job": False, + }, + "experiment": { + "base_cmd": "python run.py --a {a} --b {b}", + "a": [1, 2], + "b": ["x", "y"], + }, + }, + ) + + automation.run_experiment("exp.toml") + + assert len(fake.calls["add_task"]) == 4 + commands = [x["command_line"] for x in fake.calls["add_task"]] + assert "python run.py --a 1 --b x" in commands + assert "python run.py --a 2 --b y" in commands + assert all(x["container_image_name"] is None for x in fake.calls["add_task"]) + assert fake.calls["monitor_job"] == [] + + +def test_run_tasks_returns_none_without_pool_name(monkeypatch): + fake = FakeClient() + monkeypatch.setattr(automation, "CloudClient", lambda **kwargs: fake) + monkeypatch.setattr(automation.toml, "load", lambda _: {"job": {}, "task": []}) + + assert automation.run_tasks("tasks.toml") is None + + +def test_run_tasks_upload_dependencies_and_monitor(monkeypatch): + fake = FakeClient() + monkeypatch.setattr(automation, "CloudClient", lambda **kwargs: fake) + monkeypatch.setattr( + automation.batch_helpers, "check_pool_exists", lambda **kwargs: True + ) + monkeypatch.setattr( + automation.toml, + "load", + lambda _: { + "job": { + "pool_name": "pool-x", + "job_name": "job-3", + "container": "img:tag", + "monitor_job": True, + }, + "upload": { + "container_name": "data-cont", + "folders": ["folder-a"], + "files": ["in1.txt"], + }, + "task": [ + {"name": "prep", "cmd": "echo prep"}, + { + "name": "train", + "cmd": "echo train", + "depends_on": ["prep"], + "run_dependent_tasks_on_fail": True, + }, + ], + }, + ) + + assert automation.run_tasks("tasks.toml") is None + + assert fake.calls["upload_folders"][0]["location_in_blob"] == "" + assert fake.calls["upload_files"][0]["location_in_blob"] == "" + assert fake.calls["create_job"][0]["job_name"] == "job-3" + + first_task = fake.calls["add_task"][0] + second_task = fake.calls["add_task"][1] + assert first_task["depends_on"] is None + assert second_task["depends_on"] == ["tid-1"] + assert second_task["run_dependent_tasks_on_fail"] is True + assert second_task["container_image_name"] == "img:tag" + assert fake.calls["monitor_job"] == ["job-3"] + + +def test_run_experiment_upload_default_location_and_no_monitor_key(monkeypatch): + fake = FakeClient() + monkeypatch.setattr(automation, "CloudClient", lambda **kwargs: fake) + monkeypatch.setattr( + automation.batch_helpers, "check_pool_exists", lambda **kwargs: True + ) + monkeypatch.setattr( + automation.toml, + "load", + lambda _: { + "job": { + "pool_name": "pool-x", + "job_name": "job-4", + }, + "upload": { + "container_name": "data-cont", + "folders": ["folder-a"], + "files": ["in1.txt"], + }, + "experiment": { + "base_cmd": "echo {x}", + "x": ["ok"], + }, + }, + ) + + automation.run_experiment("exp.toml") + + assert fake.calls["upload_folders"][0]["location_in_blob"] == "" + assert fake.calls["upload_files"][0]["location_in_blob"] == "" + assert fake.calls["monitor_job"] == [] + + +def test_run_tasks_returns_none_when_client_creation_fails(monkeypatch): + monkeypatch.setattr( + automation, + "CloudClient", + lambda **kwargs: (_ for _ in ()).throw(RuntimeError("boom")), + ) + monkeypatch.setattr( + automation.toml, + "load", + lambda _: { + "job": {"pool_name": "pool", "job_name": "job"}, + "task": [{"name": "prep", "cmd": "echo prep"}], + }, + ) + + assert automation.run_tasks("tasks.toml") is None + + +def test_run_tasks_returns_none_when_pool_missing(monkeypatch): + fake = FakeClient() + monkeypatch.setattr(automation, "CloudClient", lambda **kwargs: fake) + monkeypatch.setattr( + automation.toml, + "load", + lambda _: { + "job": {"pool_name": "pool-x", "job_name": "job"}, + "task": [{"name": "prep", "cmd": "echo prep"}], + }, + ) + monkeypatch.setattr( + automation.batch_helpers, "check_pool_exists", lambda **kwargs: False + ) + + assert automation.run_tasks("tasks.toml") is None + + +def test_run_tasks_optional_job_fields_present_and_monitor_false(monkeypatch): + fake = FakeClient() + monkeypatch.setattr(automation, "CloudClient", lambda **kwargs: fake) + monkeypatch.setattr( + automation.batch_helpers, "check_pool_exists", lambda **kwargs: True + ) + monkeypatch.setattr( + automation.toml, + "load", + lambda _: { + "job": { + "pool_name": "pool-x", + "job_name": "job-5", + "save_logs_to_blob": "logs", + "logs_folder": "folder", + "task_retries": 3, + "monitor_job": False, + }, + "upload": { + "container_name": "data-cont", + "location_in_blob": "inputs", + "files": ["in1.txt"], + }, + "task": [{"name": "prep", "cmd": "echo prep"}], + }, + ) + + automation.run_tasks("tasks.toml") + + assert fake.calls["upload_files"][0]["location_in_blob"] == "inputs" + assert fake.calls["create_job"][0]["save_logs_to_blob"] == "logs" + assert fake.calls["create_job"][0]["logs_folder"] == "folder" + assert fake.calls["create_job"][0]["task_retries"] == 3 + assert fake.calls["monitor_job"] == [] + + +def test_run_tasks_no_upload_no_container_and_no_monitor_key(monkeypatch): + fake = FakeClient() + monkeypatch.setattr(automation, "CloudClient", lambda **kwargs: fake) + monkeypatch.setattr( + automation.batch_helpers, "check_pool_exists", lambda **kwargs: True + ) + monkeypatch.setattr( + automation.toml, + "load", + lambda _: { + "job": { + "pool_name": "pool-x", + "job_name": "job-6", + }, + "task": [{"name": "prep", "cmd": "echo prep"}], + }, + ) + + automation.run_tasks("tasks.toml") + + assert fake.calls["upload_folders"] == [] + assert fake.calls["upload_files"] == [] + assert fake.calls["add_task"][0]["container_image_name"] is None + assert fake.calls["monitor_job"] == [] diff --git a/tests/test_client_task_helpers.py b/tests/test_client_task_helpers.py new file mode 100644 index 0000000..b50fa7b --- /dev/null +++ b/tests/test_client_task_helpers.py @@ -0,0 +1,468 @@ +import json +import logging +from types import SimpleNamespace + +import pytest +from azure.batch.models import BatchNodeIdentityReference + +from cfa.cloudops import client, helpers, task + + +@pytest.fixture +def fake_credential_handler(): + return SimpleNamespace( + method="sp", + client_secret_credential="sp-cred", # pragma: allowlist secret + client_secret_sp_credential="default-cred", # pragma: allowlist secret + user_credential="user-cred", + azure_subscription_id="sub-123", + azure_batch_endpoint="https://batch.example", + azure_blob_storage_endpoint="https://blob.example", + ) + + +def test_get_batch_management_client_methods(monkeypatch, fake_credential_handler): + calls = [] + + def fake_constructor(**kwargs): + calls.append(kwargs) + return SimpleNamespace(kind="batch_mgmt") + + monkeypatch.setattr("cfa.cloudops.client.BatchManagementClient", fake_constructor) + + fake_credential_handler.method = "sp" + client.get_batch_management_client(fake_credential_handler) + + fake_credential_handler.method = "default" + client.get_batch_management_client(fake_credential_handler) + + fake_credential_handler.method = "user" + client.get_batch_management_client(fake_credential_handler) + + assert calls[0]["credential"] == "sp-cred" + assert calls[1]["credential"] == "default-cred" + assert calls[2]["credential"] == "user-cred" + assert all(c["subscription_id"] == "sub-123" for c in calls) + + +def test_get_compute_management_client_methods(monkeypatch, fake_credential_handler): + calls = [] + + def fake_constructor(**kwargs): + calls.append(kwargs) + return SimpleNamespace(kind="compute_mgmt") + + monkeypatch.setattr("cfa.cloudops.client.ComputeManagementClient", fake_constructor) + + fake_credential_handler.method = "sp" + client.get_compute_management_client(fake_credential_handler) + + fake_credential_handler.method = "default" + client.get_compute_management_client(fake_credential_handler) + + fake_credential_handler.method = "user" + client.get_compute_management_client(fake_credential_handler) + + assert calls[0]["credential"] == "sp-cred" + assert calls[1]["credential"] == "default-cred" + assert calls[2]["credential"] == "user-cred" + + +def test_get_batch_service_client_methods(monkeypatch, fake_credential_handler): + calls = [] + + def fake_constructor(**kwargs): + calls.append(kwargs) + return SimpleNamespace(kind="batch_service") + + monkeypatch.setattr("cfa.cloudops.client.BatchClient", fake_constructor) + + fake_credential_handler.method = "sp" + client.get_batch_service_client(fake_credential_handler) + + fake_credential_handler.method = "default" + client.get_batch_service_client(fake_credential_handler) + + fake_credential_handler.method = "user" + client.get_batch_service_client(fake_credential_handler) + + assert calls[0]["credential"] == "sp-cred" + assert calls[1]["credential"] == "default-cred" + assert calls[2]["credential"] == "user-cred" + assert all(c["endpoint"] == "https://batch.example" for c in calls) + + +def test_get_blob_service_client_methods(monkeypatch, fake_credential_handler): + calls = [] + + def fake_constructor(**kwargs): + calls.append(kwargs) + return SimpleNamespace(kind="blob_service") + + monkeypatch.setattr("cfa.cloudops.client.BlobServiceClient", fake_constructor) + + fake_credential_handler.method = "sp" + client.get_blob_service_client(fake_credential_handler) + + fake_credential_handler.method = "default" + client.get_blob_service_client(fake_credential_handler) + + fake_credential_handler.method = "user" + client.get_blob_service_client(fake_credential_handler) + + assert calls[0]["credential"] == "sp-cred" + assert calls[1]["credential"] == "default-cred" + assert calls[2]["credential"] == "user-cred" + assert all(c["account_url"] == "https://blob.example" for c in calls) + + +def test_get_clients_build_default_handler_when_none( + monkeypatch, fake_credential_handler +): + monkeypatch.setattr( + "cfa.cloudops.client.EnvCredentialHandler", lambda: fake_credential_handler + ) + monkeypatch.setattr( + "cfa.cloudops.client.BatchManagementClient", lambda **kwargs: kwargs + ) + + result = client.get_batch_management_client() + assert result["subscription_id"] == "sub-123" + + +def test_create_bind_mount_string(): + mount = task.create_bind_mount_string("/mnt/batch/tasks/fsmounts", "src", "/app") + assert mount == "--mount type=bind,source=/mnt/batch/tasks/fsmounts/src,target=/app" + + +def test_get_container_settings_with_mounts_and_registry(monkeypatch): + registry = SimpleNamespace(registry_server="myregistry.azurecr.io") + settings = task.get_container_settings( + container_image_name="myregistry.azurecr.io/app:latest", + mount_pairs=[ + {"source": "input", "target": "/app/input"}, + {"source": "output", "target": "/app/output"}, + ], + additional_options="--ipc=host", + registry=registry, + ) + + assert settings.image_name == "myregistry.azurecr.io/app:latest" + assert "--ipc=host" in settings.container_run_options + assert "source=/mnt/batch/tasks/fsmounts/input" in settings.container_run_options + assert "target=/app/output" in settings.container_run_options + + +def test_output_task_files_to_blob_uses_default_identity(monkeypatch): + mgmt_id = SimpleNamespace( + resource_id="/subscriptions/sub/resourceGroups/rg/providers/id" + ) + node_id = BatchNodeIdentityReference(resource_id=mgmt_id.resource_id) + + monkeypatch.setattr( + "cfa.cloudops.task.get_compute_node_identity_reference", lambda: mgmt_id + ) + monkeypatch.setattr("cfa.cloudops.task.get_batch_compute_id", lambda x: node_id) + + output = task.output_task_files_to_blob( + file_pattern="*.txt", + blob_container="logs", + blob_account="acct", + path="job/task", + ) + + assert output.file_pattern == "*.txt" + container = output.destination.container + assert container.path == "job/task" + assert container.identity_reference.resource_id == mgmt_id.resource_id + assert container.container_url == "https://acct.blob.core.windows.net/logs" + + +def test_output_task_files_to_blob_type_error(): + with pytest.raises(TypeError): + task.output_task_files_to_blob( + file_pattern="*.txt", + blob_container="logs", + blob_account="acct", + compute_node_identity_reference="not-a-node-id", + ) + + +def test_get_task_config_with_logs_and_filtered_kwargs(monkeypatch): + node_id = BatchNodeIdentityReference( + resource_id="/subscriptions/sub/resourceGroups/rg/providers/id" + ) + log_output_file = task.output_task_files_to_blob( + file_pattern="../std*.txt", + blob_container="logs", + blob_account="acct", + path="preexisting", + compute_node_identity_reference=node_id, + ) + + monkeypatch.setattr( + "cfa.cloudops.task.output_task_files_to_blob", + lambda **kwargs: log_output_file, + ) + + cfg = task.get_task_config( + task_id="task-001", + base_call="python app.py", + output_files=[log_output_file], + log_blob_container="logs", + log_blob_account="acct", + log_subdir="jobs/run-1", + run_dependent_tasks_on_failure=True, + ) + + assert cfg.id == "task-001" + assert cfg.command_line == "python app.py" + assert len(cfg.output_files) == 2 + assert cfg.user_identity.auto_user.elevation_level.name.lower() == "admin" + + +def test_get_batch_compute_id_validation(): + valid = SimpleNamespace( + resource_id="/subscriptions/sub/resourceGroups/rg/providers/id" + ) + result = task.get_batch_compute_id(valid) + assert isinstance(result, BatchNodeIdentityReference) + assert result.resource_id == valid.resource_id + + with pytest.raises(ValueError): + task.get_batch_compute_id(SimpleNamespace(resource_id="")) + + +def test_get_log_level_variants(monkeypatch): + monkeypatch.delenv("LOG_LEVEL", raising=False) + assert helpers.get_log_level() == logging.CRITICAL + 1 + + monkeypatch.setenv("LOG_LEVEL", "none") + assert helpers.get_log_level() == logging.CRITICAL + 1 + + monkeypatch.setenv("LOG_LEVEL", "debug") + assert helpers.get_log_level() == logging.DEBUG + + monkeypatch.setenv("LOG_LEVEL", "warn") + assert helpers.get_log_level() == logging.WARNING + + monkeypatch.setenv("LOG_LEVEL", "weird") + assert helpers.get_log_level() == logging.DEBUG + + +def test_format_rel_path(): + assert helpers.format_rel_path("/data/input") == "data/input" + assert helpers.format_rel_path("data/output") == "data/output" + + +def test_list_acr_tags_success(monkeypatch): + tags = ["latest", "v1"] + responses = [ + SimpleNamespace(returncode=0, stdout="", stderr=""), + SimpleNamespace(returncode=0, stdout=json.dumps(tags), stderr=""), + ] + + def fake_run(*args, **kwargs): + return responses.pop(0) + + monkeypatch.setattr("cfa.cloudops.helpers.sp.run", fake_run) + + result = helpers.list_acr_tags("reg", "repo") + assert result == tags + + +def test_list_acr_tags_identity_login_then_success(monkeypatch): + tags = ["latest"] + responses = [ + SimpleNamespace(returncode=1, stdout="", stderr="not logged in"), + SimpleNamespace(returncode=0, stdout="", stderr=""), + SimpleNamespace(returncode=0, stdout=json.dumps(tags), stderr=""), + ] + + def fake_run(*args, **kwargs): + return responses.pop(0) + + monkeypatch.setattr("cfa.cloudops.helpers.sp.run", fake_run) + + result = helpers.list_acr_tags("reg", "repo") + assert result == tags + + +def test_list_acr_tags_identity_fails_but_existing_session(monkeypatch): + tags = ["v2"] + responses = [ + SimpleNamespace(returncode=1, stdout="", stderr="auth missing"), + SimpleNamespace(returncode=1, stdout="", stderr="identity not available"), + SimpleNamespace(returncode=0, stdout="", stderr=""), + SimpleNamespace(returncode=0, stdout=json.dumps(tags), stderr=""), + ] + + def fake_run(*args, **kwargs): + return responses.pop(0) + + monkeypatch.setattr("cfa.cloudops.helpers.sp.run", fake_run) + + result = helpers.list_acr_tags("reg", "repo") + assert result == tags + + +def test_list_acr_tags_auth_fails_hard(monkeypatch): + responses = [ + SimpleNamespace(returncode=1, stdout="", stderr="auth missing"), + SimpleNamespace(returncode=1, stdout="", stderr="identity not available"), + SimpleNamespace(returncode=1, stdout="", stderr="still no login"), + ] + + def fake_run(*args, **kwargs): + return responses.pop(0) + + monkeypatch.setattr("cfa.cloudops.helpers.sp.run", fake_run) + + with pytest.raises(Exception): + helpers.list_acr_tags("reg", "repo") + + +def test_list_acr_tags_show_tags_failure(monkeypatch): + responses = [ + SimpleNamespace(returncode=0, stdout="", stderr=""), + SimpleNamespace(returncode=3, stdout="", stderr="acr failed"), + ] + + def fake_run(*args, **kwargs): + return responses.pop(0) + + monkeypatch.setattr("cfa.cloudops.helpers.sp.run", fake_run) + + with pytest.raises(Exception): + helpers.list_acr_tags("reg", "repo") + + +@pytest.mark.parametrize( + "use_device_code,expected_login_cmd", + [ + (False, "az login --identity"), + (True, "az login --use-device-code"), + ], +) +def test_package_and_upload_dockerfile_success( + monkeypatch, use_device_code, expected_login_cmd +): + docker_env = SimpleNamespace(ping=lambda: True) + commands = [] + + def fake_run(cmd, shell=True, **kwargs): + commands.append(cmd) + return SimpleNamespace(returncode=0, stdout="", stderr="") + + monkeypatch.setattr( + "cfa.cloudops.helpers.docker.from_env", lambda timeout=10: docker_env + ) + monkeypatch.setattr("cfa.cloudops.helpers.os.path.exists", lambda p: True) + monkeypatch.setattr("cfa.cloudops.helpers.sp.run", fake_run) + + image_name = helpers.package_and_upload_dockerfile( + registry_name="reg", + repo_name="repo", + tag="v1", + path_to_dockerfile="./Dockerfile", + use_device_code=use_device_code, + ) + + assert image_name == "reg.azurecr.io/repo:v1" + assert commands[0].startswith( + "docker image build -f ./Dockerfile -t reg.azurecr.io/repo:v1" + ) + assert expected_login_cmd in commands + assert "az acr login --name reg" in commands + assert "docker push reg.azurecr.io/repo:v1" in commands + + +def test_package_and_upload_dockerfile_docker_not_running(monkeypatch): + def boom(timeout=10): + raise helpers.DockerException("down") + + monkeypatch.setattr("cfa.cloudops.helpers.docker.from_env", boom) + + with pytest.raises(helpers.DockerException): + helpers.package_and_upload_dockerfile("reg", "repo", "latest") + + +def test_package_and_upload_dockerfile_missing_file(monkeypatch): + docker_env = SimpleNamespace(ping=lambda: True) + + monkeypatch.setattr( + "cfa.cloudops.helpers.docker.from_env", lambda timeout=10: docker_env + ) + monkeypatch.setattr("cfa.cloudops.helpers.os.path.exists", lambda p: False) + + with pytest.raises(Exception): + helpers.package_and_upload_dockerfile("reg", "repo", "latest") + + +@pytest.mark.parametrize( + "use_device_code,expected_login_cmd", + [ + (False, "az login --identity"), + (True, "az login --use-device-code"), + ], +) +def test_upload_docker_image_success(monkeypatch, use_device_code, expected_login_cmd): + tagged = [] + commands = [] + image = SimpleNamespace(tag=lambda tag_name: tagged.append(tag_name)) + images = SimpleNamespace( + get=lambda image_name: image, + list=lambda: [SimpleNamespace(tags=["local:latest"])], + ) + docker_env = SimpleNamespace(ping=lambda: True, images=images) + + def fake_run(cmd, shell=True, **kwargs): + commands.append(cmd) + return SimpleNamespace(returncode=0, stdout="", stderr="") + + monkeypatch.setattr( + "cfa.cloudops.helpers.docker.from_env", lambda timeout=8: docker_env + ) + monkeypatch.setattr("cfa.cloudops.helpers.sp.run", fake_run) + + image_name = helpers.upload_docker_image( + image_name="local:latest", + registry_name="reg", + repo_name="repo", + tag="v2", + use_device_code=use_device_code, + ) + + assert image_name == "reg.azurecr.io/repo:v2" + assert tagged == ["reg.azurecr.io/repo:v2"] + assert expected_login_cmd in commands + assert "az acr login --name reg" in commands + assert "docker push reg.azurecr.io/repo:v2" in commands + + +def test_upload_docker_image_docker_not_running(monkeypatch): + def boom(timeout=8): + raise helpers.DockerException("down") + + monkeypatch.setattr("cfa.cloudops.helpers.docker.from_env", boom) + + with pytest.raises(helpers.DockerException): + helpers.upload_docker_image("local:latest", "reg", "repo") + + +def test_upload_docker_image_not_found(monkeypatch): + def missing_image(image_name): + raise helpers.docker.errors.ImageNotFound("missing") + + images = SimpleNamespace( + get=missing_image, + list=lambda: [SimpleNamespace(tags=["other:tag"])], + ) + docker_env = SimpleNamespace(ping=lambda: True, images=images) + + monkeypatch.setattr( + "cfa.cloudops.helpers.docker.from_env", lambda timeout=8: docker_env + ) + + with pytest.raises(helpers.docker.errors.ImageNotFound): + helpers.upload_docker_image("local:latest", "reg", "repo") diff --git a/tests/test_cloudclient.py b/tests/test_cloudclient.py index 72748a9..dc853e9 100644 --- a/tests/test_cloudclient.py +++ b/tests/test_cloudclient.py @@ -416,31 +416,40 @@ def test_add_task( with patch( "cfa.cloudops.batch_helpers.get_rel_mnt_path", return_value="/mnt/logs/" ): - result = cloud_client_with_service_principal.add_task( - job_name=job_name, - command_line="echo Hello World", - container_image_name="my-image:latest", - ) - assert result is True - result = cloud_client_with_service_principal.add_task( - job_name=job_name, - command_line="echo Hello World", - ) - assert result is True - cloud_client_with_service_principal.save_logs_to_blob = True - result = cloud_client_with_service_principal.add_task( - job_name=job_name, - command_line="echo Hello World", - mount_pairs=mount_pairs, - ) - assert result is True - cloud_client_with_service_principal.full_container_name = "my-container" - result = cloud_client_with_service_principal.add_task( - job_name=job_name, - command_line="echo Hello World", - mount_pairs=mount_pairs, - ) - assert result is True + # add_task now validates that a container image can be resolved from pool + mock_pool_info = MagicMock() + ( + mock_pool_info.deployment_configuration.virtual_machine_configuration.container_configuration.container_image_names + ) = ["myregistry.azurecr.io/my-image:latest"] + with patch( + "cfa.cloudops.batch_helpers.get_pool_full_info", + return_value=mock_pool_info, + ): + result = cloud_client_with_service_principal.add_task( + job_name=job_name, + command_line="echo Hello World", + container_image_name="my-image:latest", + ) + assert result is True + result = cloud_client_with_service_principal.add_task( + job_name=job_name, + command_line="echo Hello World", + ) + assert result is True + cloud_client_with_service_principal.save_logs_to_blob = True + result = cloud_client_with_service_principal.add_task( + job_name=job_name, + command_line="echo Hello World", + mount_pairs=mount_pairs, + ) + assert result is True + cloud_client_with_service_principal.full_container_name = "my-container" + result = cloud_client_with_service_principal.add_task( + job_name=job_name, + command_line="echo Hello World", + mount_pairs=mount_pairs, + ) + assert result is True def test_add_task_uses_job_info_as_dict_for_pool_lookup( @@ -544,6 +553,7 @@ def test_list_available_images( def test_run_dag(cloud_client_with_service_principal, mock_logging): cloud_client_with_service_principal.create_job("dag_job", pool_name="test_pool") + cloud_client_with_service_principal.full_container_name = "my-image:latest" t1 = Task("python step1.py") t2 = Task("python step2.py") t3 = Task("python step3.py") diff --git a/tests/test_cloudclient_more.py b/tests/test_cloudclient_more.py new file mode 100644 index 0000000..f24f453 --- /dev/null +++ b/tests/test_cloudclient_more.py @@ -0,0 +1,326 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from cfa.cloudops._cloudclient import CloudClient + + +@pytest.fixture +def cloud_client_more(monkeypatch): + cred = SimpleNamespace( + azure_resource_group_name="rg", + azure_batch_account="acct", + azure_blob_storage_account="blobacct", + azure_blob_storage_endpoint="https://blobacct.blob.core.windows.net/", + user_credential="user-cred", + client_secret_sp_credential="default-cred", # pragma: allowlist secret + client_secret_credential="sp-cred", # pragma: allowlist secret + compute_node_identity_reference=SimpleNamespace(resource_id="rid"), + ) + + with ( + patch("cfa.cloudops._cloudclient.EnvCredentialHandler", return_value=cred), + patch( + "cfa.cloudops._cloudclient.get_batch_management_client", + return_value=MagicMock(), + ), + patch( + "cfa.cloudops._cloudclient.get_compute_management_client", + return_value=MagicMock(), + ), + patch( + "cfa.cloudops._cloudclient.get_batch_service_client", + return_value=MagicMock(), + ), + patch( + "cfa.cloudops._cloudclient.get_blob_service_client", + return_value=MagicMock(), + ), + ): + return CloudClient(dotenv_path=None, use_sp=False, use_federated=False) + + +def test_check_credentials_env_default_sp(cloud_client_more, monkeypatch): + seen_creds = [] + + class FakeSub: + def __init__(self): + self.subscription_id = "sub-1" + self.display_name = "sub-name" + self.state = "Enabled" + + def fake_subscription_client(cred): + seen_creds.append(cred) + return SimpleNamespace(subscriptions=SimpleNamespace(list=lambda: [FakeSub()])) + + monkeypatch.setattr( + "cfa.cloudops._cloudclient.SubscriptionClient", fake_subscription_client + ) + + cloud_client_more.method = "env" + cloud_client_more.check_credentials() + + cloud_client_more.method = "default" + cloud_client_more.check_credentials() + + cloud_client_more.method = "sp" + cloud_client_more.check_credentials() + + assert seen_creds == ["user-cred", "default-cred", "sp-cred"] + + +def test_check_credentials_handles_exception(cloud_client_more, monkeypatch): + monkeypatch.setattr( + "cfa.cloudops._cloudclient.SubscriptionClient", + lambda cred: (_ for _ in ()).throw(RuntimeError("boom")), + ) + + cloud_client_more.check_credentials() + + +def test_delete_job_and_schedule_methods(cloud_client_more): + cloud_client_more.batch_service_client.begin_delete_job.return_value.result.return_value = None + cloud_client_more.batch_service_client.begin_delete_job_schedule.return_value.result.return_value = None + + cloud_client_more.delete_job("job-1") + cloud_client_more.delete_job_schedule("sched-1") + cloud_client_more.resume_job_schedule("sched-1") + cloud_client_more.suspend_job_schedule("sched-1") + + cloud_client_more.batch_service_client.begin_delete_job.assert_called_once_with( + "job-1" + ) + cloud_client_more.batch_service_client.begin_delete_job_schedule.assert_called_once_with( + "sched-1" + ) + cloud_client_more.batch_service_client.enable_job_schedule.assert_called_once_with( + "sched-1" + ) + cloud_client_more.batch_service_client.disable_job_schedule.assert_called_once_with( + "sched-1" + ) + + +def test_list_available_images_filtering(cloud_client_more): + linux = SimpleNamespace(os_type="linux", name="lin") + windows = SimpleNamespace(os_type="windows", name="win") + + cloud_client_more.batch_service_client.list_supported_images.return_value = [ + linux, + windows, + ] + + with patch("cfa.cloudops._cloudclient.batch_models.OSType") as os_type: + os_type.linux = "linux" + os_type.windows = "windows" + + result_linux = cloud_client_more.list_available_images("linux") + result_windows = cloud_client_more.list_available_images("windows") + result_all = cloud_client_more.list_available_images() + + assert result_linux == [linux] + assert result_windows == [windows] + assert result_all == [linux, windows] + + +def test_package_and_upload_dockerfile_delegates(cloud_client_more, monkeypatch): + monkeypatch.setattr( + "cfa.cloudops._cloudclient.helpers.package_and_upload_dockerfile", + lambda *args, **kwargs: "reg.azurecr.io/repo:v1", + ) + + out = cloud_client_more.package_and_upload_dockerfile("reg", "repo", "v1") + + assert out == "reg.azurecr.io/repo:v1" + assert cloud_client_more.container_registry_server == "reg.azurecr.io" + assert cloud_client_more.registry_url == "https://reg.azurecr.io" + + +def test_upload_docker_image_delegates(cloud_client_more, monkeypatch): + monkeypatch.setattr( + "cfa.cloudops._cloudclient.helpers.upload_docker_image", + lambda *args, **kwargs: "reg.azurecr.io/repo:v2", + ) + + out = cloud_client_more.upload_docker_image("local:latest", "reg", "repo", "v2") + + assert out == "reg.azurecr.io/repo:v2" + assert cloud_client_more.container_image_name == "https://reg.azurecr.io/repo:v2" + + +def test_download_file_download_folder_delegates(cloud_client_more, monkeypatch): + seen = {"file": None, "folder": None} + + monkeypatch.setattr( + "cfa.cloudops._cloudclient.blob_helpers.download_file", + lambda *args, **kwargs: seen.__setitem__("file", (args, kwargs)), + ) + monkeypatch.setattr( + "cfa.cloudops._cloudclient.blob_helpers.download_folder", + lambda *args, **kwargs: seen.__setitem__("folder", (args, kwargs)), + ) + + cloud_client_more.download_file("a.txt", "./a.txt", container_name="c1") + cloud_client_more.download_folder("src", "dest", "c2") + + assert seen["file"] is not None + assert seen["folder"] is not None + + +def test_async_download_folder_uses_method_credential(cloud_client_more, monkeypatch): + calls = [] + + def fake_async_download_blob_folder(**kwargs): + calls.append(kwargs) + + monkeypatch.setattr( + "cfa.cloudops._cloudclient.blob.async_download_blob_folder", + fake_async_download_blob_folder, + ) + + cloud_client_more.method = "default" + cloud_client_more.async_download_folder("src", "dest", "c") + + cloud_client_more.method = "sp" + cloud_client_more.async_download_folder("src", "dest", "c") + + cloud_client_more.method = "env" + cloud_client_more.async_download_folder("src", "dest", "c") + + assert calls[0]["credential"] == "default-cred" + assert calls[1]["credential"] == "sp-cred" + assert calls[2]["credential"] == "user-cred" + + +def test_async_upload_folder_handles_str_and_list(cloud_client_more, monkeypatch): + calls = [] + + monkeypatch.setattr( + "cfa.cloudops._cloudclient.blob.async_upload_folder", + lambda **kwargs: calls.append(kwargs), + ) + + cloud_client_more.method = "env" + cloud_client_more.async_upload_folder("folder-a", "container-a") + cloud_client_more.async_upload_folder(["folder-b", "folder-c"], "container-a") + + assert [c["folder"] for c in calls] == ["folder-a", "folder-b", "folder-c"] + + +def test_delete_pool_and_blob_ops(cloud_client_more, monkeypatch): + called = {"pool": 0, "blob_file": 0, "blob_folder": 0} + + monkeypatch.setattr( + "cfa.cloudops._cloudclient.batch_helpers.delete_pool", + lambda **kwargs: called.__setitem__("pool", called["pool"] + 1), + ) + monkeypatch.setattr( + "cfa.cloudops._cloudclient.blob_helpers.delete_blob_snapshots", + lambda *args, **kwargs: called.__setitem__( + "blob_file", called["blob_file"] + 1 + ), + ) + monkeypatch.setattr( + "cfa.cloudops._cloudclient.blob_helpers.delete_blob_folder", + lambda *args, **kwargs: called.__setitem__( + "blob_folder", called["blob_folder"] + 1 + ), + ) + + cloud_client_more.delete_pool("pool-1") + cloud_client_more.delete_blob_file("blob.txt", "cont") + cloud_client_more.delete_blob_folder("folder", "cont") + + assert called == {"pool": 1, "blob_file": 1, "blob_folder": 1} + + +def test_list_blob_files_by_container_and_mounts(cloud_client_more, monkeypatch): + monkeypatch.setattr( + "cfa.cloudops._cloudclient.blob_helpers.list_blobs_flat", + lambda container_name, blob_service_client, verbose=False: [ + f"{container_name}/a.txt" + ], + ) + + out_container = cloud_client_more.list_blob_files("c1") + + cloud_client_more.mounts = [("m1", "m1"), ("m2", "m2")] + out_mounts = cloud_client_more.list_blob_files() + + assert out_container == ["c1/a.txt"] + assert out_mounts == ["m1/a.txt", "m2/a.txt"] + + +def test_download_job_stats_and_task_status(cloud_client_more, monkeypatch): + seen = {"stats": None, "status": None} + + monkeypatch.setattr( + "cfa.cloudops._cloudclient.batch_helpers.download_job_stats", + lambda **kwargs: seen.__setitem__("stats", kwargs), + ) + monkeypatch.setattr( + "cfa.cloudops._cloudclient.batch_helpers.get_task_status", + lambda **kwargs: seen.__setitem__("status", kwargs) or "{}", + ) + + cloud_client_more.download_job_stats("job-1") + status = cloud_client_more.get_task_status("job-1", "task-1") + + assert seen["stats"]["job_name"] == "job-1" + assert status == "{}" + + +def test_download_after_job_dispatches_file_and_folder(cloud_client_more, monkeypatch): + calls = {"file": [], "folder": [], "monitor": 0, "makedirs": 0} + + monkeypatch.setattr( + "cfa.cloudops._cloudclient.batch_helpers.monitor_tasks", + lambda **kwargs: calls.__setitem__("monitor", calls["monitor"] + 1), + ) + monkeypatch.setattr( + "cfa.cloudops._cloudclient.os.makedirs", + lambda *args, **kwargs: calls.__setitem__("makedirs", calls["makedirs"] + 1), + ) + + cloud_client_more.download_file = lambda **kwargs: calls["file"].append(kwargs) + cloud_client_more.download_folder = lambda **kwargs: calls["folder"].append(kwargs) + + cloud_client_more.download_after_job( + job_name="job-1", + blob_paths=["outputs/file.txt", "outputs/folder"], + target="./downloads", + container_name="cont", + ) + + assert calls["monitor"] == 1 + assert calls["makedirs"] == 1 + assert len(calls["file"]) == 1 + assert len(calls["folder"]) == 1 + + +def test_get_kv_secret_success_and_failure(cloud_client_more, monkeypatch): + monkeypatch.setattr( + "cfa.cloudops._cloudclient.SecretClient", + lambda vault_url, credential: SimpleNamespace( + get_secret=lambda name: SimpleNamespace(value=f"v-{name}") + ), + ) + + cloud_client_more.method = "sp" + assert cloud_client_more.get_kv_secret("s1", "kv") == "v-s1" + + monkeypatch.setattr( + "cfa.cloudops._cloudclient.SecretClient", + lambda vault_url, credential: (_ for _ in ()).throw(RuntimeError("boom")), + ) + assert cloud_client_more.get_kv_secret("s1", "kv") is None + + +def test_list_acr_tags_delegates(cloud_client_more, monkeypatch): + monkeypatch.setattr( + "cfa.cloudops._cloudclient.helpers.list_acr_tags", + lambda registry_name, repo_name: ["latest", "v1"], + ) + + assert cloud_client_more.list_acr_tags("reg", "repo") == ["latest", "v1"] diff --git a/tests/test_config_endpoints_defaults_util.py b/tests/test_config_endpoints_defaults_util.py new file mode 100644 index 0000000..4ef1e58 --- /dev/null +++ b/tests/test_config_endpoints_defaults_util.py @@ -0,0 +1,253 @@ +import os +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from azure.mgmt.batch import models + +from cfa.cloudops import config, defaults, endpoints, util + + +def test_try_get_val_from_dict_success_and_missing(): + value, message = config.try_get_val_from_dict("my_key", {"my_key": "value"}) + assert value == "value" + assert message is None + + value, message = config.try_get_val_from_dict("missing", {"my_key": "value"}) + assert value is None + assert "missing" in message + + +def test_try_get_val_from_env_success_and_missing(monkeypatch): + monkeypatch.setenv("MY_ENV", "abc") + value, message = config.try_get_val_from_env("MY_ENV") + assert value == "abc" + assert message is None + + monkeypatch.delenv("MISSING_ENV", raising=False) + value, message = config.try_get_val_from_env("MISSING_ENV") + assert value is None + assert "MISSING_ENV" in message + + +def test_get_config_val_prefers_dict_over_env(monkeypatch): + monkeypatch.setenv("API_KEY", "env-value") + result = config.get_config_val("api_key", config_dict={"api_key": "dict-value"}) + assert result == "dict-value" + + +def test_get_config_val_falls_back_to_env(monkeypatch): + monkeypatch.setenv("API_KEY", "env-value") + result = config.get_config_val("api_key", config_dict={}) + assert result == "env-value" + + +def test_get_config_val_missing_returns_none(monkeypatch): + monkeypatch.delenv("NOT_FOUND", raising=False) + result = config.get_config_val( + "not_found", + config_dict={}, + try_env=True, + env_variable_name="NOT_FOUND", + ) + assert result is None + + +def test_construct_https_url(): + assert ( + endpoints._construct_https_url("example.com", "/v1") == "https://example.com/v1" + ) + + +def test_batch_blob_and_registry_endpoint_constructors(): + assert ( + endpoints.construct_batch_endpoint("acct", "eastus") + == "https://acct.eastus.batch.azure.com/" + ) + assert ( + endpoints.construct_batch_endpoint("acct", "westus", "custom.domain/") + == "https://acct.westus.custom.domain/" + ) + + assert ( + endpoints.construct_azure_container_registry_endpoint("myregistry") + == "https://myregistry.azurecr.io" + ) + assert ( + endpoints.construct_blob_account_endpoint("storage") + == "https://storage.blob.core.windows.net/" + ) + assert ( + endpoints.construct_blob_container_endpoint("my folder", "storage") + == "https://storage.blob.core.windows.net/my%20folder" + ) + + +@pytest.mark.parametrize( + "endpoint,expected_valid,expected_substring", + [ + ("https://myregistry.azurecr.io", True, None), + ("https://myregistry.azurecr.io/", False, "trailing slash"), + ("https://myregistry.example.com", False, "azurecr.io"), + ("https://azurecr.io", False, "subdomain"), + ], +) +def test_is_valid_acr_endpoint(endpoint, expected_valid, expected_substring): + is_valid, error_message = endpoints.is_valid_acr_endpoint(endpoint) + assert is_valid is expected_valid + if expected_substring is None: + assert error_message is None + else: + assert expected_substring in error_message + + +def test_remaining_task_autoscale_formula_contains_parameters(): + formula = defaults.remaining_task_autoscale_formula( + task_sample_interval_minutes=30, + max_number_vms=22, + ) + assert "TimeInterval_Minute * 30" in formula + assert "cappedPoolSize = 22" in formula + + +def test_set_env_vars_sets_defaults_and_derived_values(monkeypatch): + monkeypatch.setenv("AZURE_BATCH_ACCOUNT", "acct") + monkeypatch.setenv("AZURE_BATCH_LOCATION", "eastus") + monkeypatch.setenv("AZURE_KEYVAULT_NAME", "myvault") + monkeypatch.setenv("AZURE_BLOB_STORAGE_ACCOUNT", "blobacct") + monkeypatch.setenv("AZURE_CONTAINER_REGISTRY_ACCOUNT", "regacct") + + defaults.set_env_vars() + + assert "https://acct.eastus.batch.azure.com/" == os.environ["AZURE_BATCH_ENDPOINT"] + assert "https://myvault.vault.azure.net" == os.environ["AZURE_KEYVAULT_ENDPOINT"] + assert ( + "https://blobacct.blob.core.windows.net/" + == os.environ["AZURE_BLOB_STORAGE_ENDPOINT"] + ) + assert "regacct.azurecr.io/" == os.environ["ACR_TAG_PREFIX"] + + +def test_get_default_pool_identity_and_pool_config(): + identity_path = "/subscriptions/sub/resourceGroups/rg/providers/id" + identity = defaults.get_default_pool_identity(identity_path) + + assert identity.type == models.PoolIdentityType.user_assigned + assert identity_path in identity.user_assigned_identities + + pool = defaults.get_default_pool_config( + pool_name="pool-a", + subnet_id="/subscriptions/sub/resourceGroups/rg/providers/net/subnets/default", + user_assigned_identity=identity_path, + vm_size="standard_d2s_v3", + ) + + assert pool.display_name == "pool-a" + assert pool.vm_size == "standard_d2s_v3" + assert pool.network_configuration.subnet_id.endswith("/subnets/default") + + +def test_assign_container_config_updates_pool_in_place(): + identity_path = "/subscriptions/sub/resourceGroups/rg/providers/id" + pool = defaults.get_default_pool_config( + pool_name="pool-b", + subnet_id="/subscriptions/sub/resourceGroups/rg/providers/net/subnets/default", + user_assigned_identity=identity_path, + ) + + container_config = models.ContainerConfiguration(type="dockerCompatible") + updated = defaults.assign_container_config(pool, container_config) + + assert updated is pool + vm_config = updated.deployment_configuration.virtual_machine_configuration + assert vm_config.container_configuration is container_config + + +def test_ensure_listlike_behaviors(): + data = ["a", "b"] + assert util.ensure_listlike(data) is data + assert util.ensure_listlike("a") == ["a"] + assert util.ensure_listlike(5) == [5] + + +def test_sku_to_dict_handles_capabilities_and_properties(): + sku = SimpleNamespace( + name="Standard_D2s_v3", + family_name="standardDSv3Family", + batch_support_end_of_life="2027-01-01", + additional_properties={"tier": "standard"}, + capabilities=[SimpleNamespace(name="vCPUs", value="2")], + ) + + as_dict = util.sku_to_dict(sku) + + assert as_dict["name"] == "Standard_D2s_v3" + assert as_dict["family_name"] == "standardDSv3Family" + assert as_dict["vCPUs"] == "2" + assert as_dict["additional_properties"]["tier"] == "standard" + + +def test_get_subscriptions_success_and_failure(monkeypatch): + class FakeSub: + def __init__(self, display_name): + self.display_name = display_name + + fake_client = MagicMock() + fake_client.subscriptions.list.return_value = [FakeSub("sub-a"), FakeSub("sub-b")] + + monkeypatch.setattr("cfa.cloudops.util.DefaultAzureCredential", lambda: object()) + monkeypatch.setattr( + "cfa.cloudops.util.SubscriptionClient", lambda cred: fake_client + ) + + assert util.get_subscriptions() == ["sub-a", "sub-b"] + + monkeypatch.setattr( + "cfa.cloudops.util.DefaultAzureCredential", + lambda: (_ for _ in ()).throw(RuntimeError("boom")), + ) + assert util.get_subscriptions() == [] + + +def test_check_ext_env(monkeypatch): + monkeypatch.setattr("cfa.cloudops.util.get_subscriptions", lambda: ["foo"]) + assert util.check_ext_env() is False + + monkeypatch.setattr( + "cfa.cloudops.util.get_subscriptions", lambda: ["EXT-EDAV-CFA sandbox"] + ) + assert util.check_ext_env() is True + + +def test_get_user_fallbacks(monkeypatch): + monkeypatch.setattr("cfa.cloudops.util.getpass.getuser", lambda: "alice") + assert util.get_user() == "alice" + + monkeypatch.setattr( + "cfa.cloudops.util.getpass.getuser", + lambda: (_ for _ in ()).throw(RuntimeError("no getuser")), + ) + monkeypatch.setattr( + "cfa.cloudops.util.sp.run", + lambda *args, **kwargs: SimpleNamespace(stdout="bob\n"), + ) + assert util.get_user() == "bob" + + monkeypatch.setattr( + "cfa.cloudops.util.sp.run", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("no whoami")), + ) + assert util.get_user() == "unknown_user" + + +def test_get_date_time_success_and_failure(monkeypatch): + timestamp = util.get_date_time() + assert "T" in timestamp + + class BrokenDateTime: + @staticmethod + def now(*args, **kwargs): + raise RuntimeError("bad clock") + + monkeypatch.setattr("cfa.cloudops.util.datetime.datetime", BrokenDateTime) + assert util.get_date_time() == "unknown_datetime" diff --git a/tests/test_function_app_client_branches.py b/tests/test_function_app_client_branches.py new file mode 100644 index 0000000..467b4f6 --- /dev/null +++ b/tests/test_function_app_client_branches.py @@ -0,0 +1,295 @@ +import subprocess +from pathlib import Path +from types import SimpleNamespace + +import pandas as pd +import pytest + +from cfa.cloudops import _function_app_client as func_mod + + +@pytest.fixture +def base_client(): + c = func_mod.FunctionAppClient.__new__(func_mod.FunctionAppClient) + c.function_app_name = "fa-app" + c.cred = SimpleNamespace( + azure_resource_group_name="rg", + azure_subscription_id="sub", + azure_tenant_id="tenant", + azure_client_id="cid", + azure_client_secret="secret", # pragma: allowlist secret + azure_blob_storage_account="blob", + client_secret_credential="cred", # pragma: allowlist secret + ) + c.update_function_database = True + c.conn = None + return c + + +def test_clone_deployment_slot_success_and_failure(monkeypatch, base_client): + calls = [] + + monkeypatch.setattr( + func_mod.subprocess, "run", lambda args, check=True: calls.append(args) + ) + assert base_client._clone_deployment_slot("newslot", "rollback") is True + assert "--configuration-source" in calls[0] + assert "fa-app/rollback" in calls[0] + + def fail(*args, **kwargs): + raise subprocess.CalledProcessError(1, "az") + + monkeypatch.setattr(func_mod.subprocess, "run", fail) + assert base_client._clone_deployment_slot("newslot", "production") is False + + +def test_swap_and_delete_deployment_slot(monkeypatch, base_client): + web_apps = SimpleNamespace( + begin_swap_slot_with_production=lambda **k: SimpleNamespace( + result=lambda: "prod" + ), + begin_swap_slot=lambda **k: SimpleNamespace(result=lambda: "slot"), + delete_slot=lambda **k: "deleted", + ) + monkeypatch.setattr( + func_mod, + "WebSiteManagementClient", + lambda cred, sub: SimpleNamespace(web_apps=web_apps), + ) + + assert base_client._swap_deployment_slot("staging", "production") == "prod" + assert base_client._swap_deployment_slot("staging", "blue") == "slot" + base_client._delete_deployment_slot("rollback") + + +def test_find_available_and_allocate_function_app(monkeypatch, base_client): + class FakeSqlResult: + def __init__(self, df): + self._df = df + + def fetchdf(self): + return self._df + + class FakeConn: + def __init__(self): + self.exec_calls = [] + + def sql(self, query): + if "SELECT FunctionAppName" in query: + return FakeSqlResult(pd.DataFrame([{"FunctionAppName": "fa-01"}])) + return self + + def execute(self, query): + self.exec_calls.append(query) + return self + + fc = FakeConn() + monkeypatch.setattr(base_client, "_get_database_connection", lambda: fc) + + assert base_client._find_available_function_app() == "fa-01" + assert base_client._allocate_function_app() is True + assert any("UPDATE function_apps" in q for q in fc.exec_calls) + assert any("COPY function_apps" in q for q in fc.exec_calls) + + +def test_log_into_portal_and_restart_paths(monkeypatch, base_client): + calls = [] + monkeypatch.setattr(func_mod.time, "sleep", lambda _: None) + monkeypatch.setattr( + func_mod.subprocess, "run", lambda args, check=True: calls.append(args) + ) + + assert base_client._log_into_portal() is True + assert calls[0][0:3] == ["az", "login", "--service-principal"] + assert calls[1][0:3] == ["az", "account", "set"] + + assert base_client._restart_function() is True + + def fail(*args, **kwargs): + raise subprocess.CalledProcessError(1, "az") + + monkeypatch.setattr(func_mod.subprocess, "run", fail) + assert base_client._log_into_portal() is False + assert base_client._restart_function() is False + + +def test_enable_health_check_and_update_settings(monkeypatch, base_client): + seen = [] + monkeypatch.setattr( + func_mod.subprocess, "run", lambda args, check=True: seen.append(args) + ) + + assert base_client._enable_health_check(slot="staging") is True + assert "--slot" in seen[0] + + assert ( + base_client._update_app_settings([("A", "1"), ("B", "2")], slot="staging") + is True + ) + assert any("A=1" in x for x in seen[1]) + + def fail(*args, **kwargs): + raise subprocess.CalledProcessError(1, "az") + + monkeypatch.setattr(func_mod.subprocess, "run", fail) + assert base_client._enable_health_check() is False + assert base_client._update_app_settings([("A", "1")]) is False + + +def test_add_user_package_delete_folder_copy_template( + tmp_path, monkeypatch, base_client +): + monkeypatch.chdir(tmp_path) + + # Current implementation expects a callable name even for string input. + with pytest.raises(AttributeError): + base_client._add_user_package_to_deployment("print('x')") + + # add_user_package with callable + def user_package_func(): + return 1 + + base_client._add_user_package_to_deployment(user_package_func) + out = Path("user_package.py").read_text() + assert "def user_package_func" in out + assert "user_package_func()" in out + + # delete_deployment_folder both branches + dep = tmp_path / base_client.function_app_name + dep.mkdir() + assert base_client._delete_deployment_folder() is True + assert base_client._delete_deployment_folder() is False + + # copy_template_to_deployment + template = tmp_path / "template" + template.mkdir() + for n in [ + "timer_blueprint", + "function_app", + "containers", + "cfa_service", + "user_package", + ]: + (template / f"{n}.txt").write_text(n) + for n in ["host", "local.settings"]: + (template / f"{n}.txt").write_text(n) + (template / "requirements.txt").write_text("pytest") + + (tmp_path / base_client.function_app_name).mkdir() + base_client._copy_template_to_deployment(str(tmp_path)) + + assert (tmp_path / base_client.function_app_name / "timer_blueprint.py").exists() + assert (tmp_path / base_client.function_app_name / "host.json").exists() + assert (tmp_path / base_client.function_app_name / "requirements.txt").exists() + + +def test_publish_function_success_and_failure(tmp_path, monkeypatch, base_client): + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(func_mod.time, "sleep", lambda _: None) + + clone_calls = [] + delete_calls = [] + settings_calls = [] + + monkeypatch.setattr(base_client, "_delete_deployment_folder", lambda: True) + + def copy_template(parent_folder): + Path(base_client.function_app_name).mkdir(exist_ok=True) + Path(base_client.function_app_name, "requirements.txt").write_text("base\n") + + monkeypatch.setattr(base_client, "_copy_template_to_deployment", copy_template) + monkeypatch.setattr( + base_client, "_add_user_package_to_deployment", lambda user_package: None + ) + monkeypatch.setattr( + base_client, + "_clone_deployment_slot", + lambda slot_name, source_slot=None: ( + clone_calls.append((slot_name, source_slot)) or True + ), + ) + monkeypatch.setattr( + base_client, "_delete_deployment_slot", lambda slot: delete_calls.append(slot) + ) + monkeypatch.setattr( + base_client, + "_update_app_settings", + lambda settings, slot=None: settings_calls.append(settings) or True, + ) + monkeypatch.setattr(base_client, "_enable_health_check", lambda slot=None: True) + monkeypatch.setattr( + base_client, "_swap_deployment_slot", lambda source_slot, target_slot: True + ) + + monkeypatch.setattr(func_mod.subprocess, "run", lambda args, check=True: None) + monkeypatch.setattr( + func_mod.FunctionAppClient, "get_health_check_flag", lambda *a, **k: True + ) + monkeypatch.setattr( + func_mod.FunctionAppClient, + "list_slots", + lambda *a, **k: [ + ("fa-app/rollback", "Running", True, None, None), + ("fa-app/backup", "Running", True, None, None), + ("fa-app/rollbackprevious", "Running", True, None, None), + ], + ) + + def user_package(): + return 42 + + assert ( + base_client._publish_function( + schedule="* * * * * *", + user_package=user_package, + dependencies=["numpy==1.0"], + environment_variables=[("K", "V")], + ) + is True + ) + assert ("rollbackprevious", "rollback") in clone_calls + assert ("rollback", None) in clone_calls + assert ("backup", None) in clone_calls + assert "rollback" in delete_calls + assert "backup" in delete_calls + assert "rollbackprevious" in delete_calls + assert len(settings_calls) == 2 + + base_client.function_app_name = "fa-app-2" + monkeypatch.setattr( + func_mod.FunctionAppClient, "get_health_check_flag", lambda *a, **k: False + ) + monkeypatch.setattr( + func_mod.subprocess, + "run", + lambda *a, **k: (_ for _ in ()).throw(subprocess.CalledProcessError(1, "func")), + ) + assert base_client._publish_function("* * * * * *", user_package) is False + + +def test_deploy_function_branches(monkeypatch, base_client): + # Login failure + monkeypatch.setattr(base_client, "_log_into_portal", lambda: False) + assert base_client.deploy_function("* * * * * *", lambda: 1) is False + + # Missing function app name and no available app + base_client.function_app_name = None + monkeypatch.setattr(base_client, "_log_into_portal", lambda: True) + monkeypatch.setattr(base_client, "_find_available_function_app", lambda: None) + assert base_client.deploy_function("* * * * * *", lambda: 1) is False + + # Publish failure + monkeypatch.setattr(base_client, "_find_available_function_app", lambda: "fa-a") + monkeypatch.setattr(base_client, "_publish_function", lambda *a, **k: False) + assert base_client.deploy_function("* * * * * *", lambda: 1) is False + + # Allocate false but continue, restart fails + monkeypatch.setattr(base_client, "_publish_function", lambda *a, **k: True) + monkeypatch.setattr(base_client, "_allocate_function_app", lambda: False) + monkeypatch.setattr(base_client, "_restart_function", lambda: False) + assert base_client.deploy_function("* * * * * *", lambda: 1) is False + + # Happy path + monkeypatch.setattr(base_client, "_allocate_function_app", lambda: True) + monkeypatch.setattr(base_client, "_restart_function", lambda: True) + assert base_client.deploy_function("* * * * * *", lambda: 1) is True diff --git a/tests/test_init_module.py b/tests/test_init_module.py new file mode 100644 index 0000000..ad37577 --- /dev/null +++ b/tests/test_init_module.py @@ -0,0 +1,62 @@ +import importlib +from types import SimpleNamespace + +import pytest + + +def _reload_cloudops(monkeypatch, log_output=None): + if log_output is None: + monkeypatch.delenv("LOG_OUTPUT", raising=False) + else: + monkeypatch.setenv("LOG_OUTPUT", log_output) + + monkeypatch.setattr("logging.basicConfig", lambda **kwargs: None) + monkeypatch.setattr( + "logging.StreamHandler", lambda stream=None: SimpleNamespace(kind="stream") + ) + monkeypatch.setattr( + "logging.FileHandler", lambda path: SimpleNamespace(kind="file", path=path) + ) + + import cfa.cloudops as cloudops + + return importlib.reload(cloudops) + + +def test_cloudops_getattr_known_symbols(monkeypatch): + cloudops = _reload_cloudops(monkeypatch, None) + + from cfa.cloudops._cloudclient import CloudClient + + assert cloudops.__getattr__("CloudClient") is CloudClient + assert "CloudClient" in cloudops.__all__ + + +def test_cloudops_getattr_unknown_symbol(monkeypatch): + cloudops = _reload_cloudops(monkeypatch, None) + + with pytest.raises(AttributeError): + cloudops.__getattr__("DoesNotExist") + + +@pytest.mark.parametrize("log_output", ["both", "file", "std", "stdout"]) +def test_cloudops_log_output_variants(monkeypatch, log_output): + made_dirs = [] + + monkeypatch.setattr("os.path.exists", lambda path: False) + monkeypatch.setattr("os.mkdir", lambda path: made_dirs.append(path)) + + cloudops = _reload_cloudops(monkeypatch, log_output) + + assert cloudops is not None + if log_output.startswith("both") or log_output.startswith("file"): + assert "logs" in made_dirs + + +def test_cloudops_log_output_unrecognized(monkeypatch, capsys): + monkeypatch.setattr("os.path.exists", lambda path: True) + cloudops = _reload_cloudops(monkeypatch, "weird-output") + + assert cloudops is not None + captured = capsys.readouterr() + assert "Did not recognize weird-output" in captured.out diff --git a/tests/test_local_automation_and_metaflow_more.py b/tests/test_local_automation_and_metaflow_more.py new file mode 100644 index 0000000..a88b3f9 --- /dev/null +++ b/tests/test_local_automation_and_metaflow_more.py @@ -0,0 +1,426 @@ +import importlib +import sys +from types import ModuleType, SimpleNamespace + +import pytest + +from cfa.cloudops.local import automation as local_automation + + +@pytest.fixture +def svc_mod(monkeypatch): + fake_examples = { + "examples": ModuleType("examples"), + "examples.metaflow": ModuleType("examples.metaflow"), + "examples.metaflow.azure_batch_decorator": ModuleType( + "examples.metaflow.azure_batch_decorator" + ), + "examples.metaflow.plugins": ModuleType("examples.metaflow.plugins"), + "examples.metaflow.plugins.metadata_providers": ModuleType( + "examples.metaflow.plugins.metadata_providers" + ), + "examples.metaflow.plugins.metadata_providers.local": ModuleType( + "examples.metaflow.plugins.metadata_providers.local" + ), + } + + class FakeAzureBatchDecorator: + pass + + class FakeLocalMetadataProvider: + pass + + fake_examples[ + "examples.metaflow.azure_batch_decorator" + ].AzureBatchDecorator = FakeAzureBatchDecorator + fake_examples[ + "examples.metaflow.plugins.metadata_providers.local" + ].LocalMetadataProvider = FakeLocalMetadataProvider + + for name, mod in fake_examples.items(): + monkeypatch.setitem(sys.modules, name, mod) + + return importlib.import_module( + "cfa.cloudops.metaflow.custom_metaflow.cfa_batch_pool_service" + ) + + +class FakeLocalClient: + def __init__(self): + self.calls = { + "upload_folders": [], + "upload_files": [], + "create_job": [], + "add_task": [], + "add_tasks_from_yaml": [], + "monitor_job": [], + } + self.cont_name = None + + def upload_folders(self, **kwargs): + self.calls["upload_folders"].append(kwargs) + + def upload_files(self, **kwargs): + self.calls["upload_files"].append(kwargs) + + def create_job(self, **kwargs): + self.calls["create_job"].append(kwargs) + + def add_task(self, **kwargs): + self.calls["add_task"].append(kwargs) + return f"tid-{len(self.calls['add_task'])}" + + def add_tasks_from_yaml(self, **kwargs): + self.calls["add_tasks_from_yaml"].append(kwargs) + + def monitor_job(self, job_name): + self.calls["monitor_job"].append(job_name) + + +def test_local_run_experiment_client_creation_failure(monkeypatch): + monkeypatch.setattr( + local_automation, + "CloudClient", + lambda **kwargs: (_ for _ in ()).throw(RuntimeError("boom")), + ) + monkeypatch.setattr( + local_automation.toml, + "load", + lambda _: { + "job": {"pool_name": "p1", "job_name": "j1", "container": "img:1"}, + "experiment": {"base_cmd": "echo {x}", "x": [1]}, + }, + ) + + assert local_automation.run_experiment("exp.toml") is None + + +def test_local_run_experiment_pool_missing_file_uses_job_container(monkeypatch): + fake = FakeLocalClient() + monkeypatch.setattr(local_automation, "CloudClient", lambda **kwargs: fake) + monkeypatch.setattr(local_automation.os.path, "exists", lambda p: False) + monkeypatch.setattr( + local_automation.toml, + "load", + lambda _: { + "job": { + "pool_name": "p1", + "job_name": "j1", + "container": "repo/image:tag", + "monitor_job": True, + }, + "experiment": {"base_cmd": "echo {x}", "x": [1]}, + }, + ) + + local_automation.run_experiment("exp.toml") + + assert fake.calls["create_job"][0]["pool_name"] == "p1" + assert fake.calls["add_task"][0]["container_image_name"] == "repo_image_tag.j1" + assert fake.calls["monitor_job"] == ["j1"] + + +def test_local_run_experiment_pool_exists_docker_errors_and_yaml(monkeypatch): + fake = FakeLocalClient() + monkeypatch.setattr(local_automation, "CloudClient", lambda **kwargs: fake) + monkeypatch.setattr(local_automation.os.path, "exists", lambda p: True) + monkeypatch.setattr( + local_automation.Path, + "read_text", + lambda self: "{'image_name': 'repo/image:tag'}", + ) + + class FakeDockerEnv: + def ping(self): + raise RuntimeError("docker down") + + class images: + @staticmethod + def get(name): + raise RuntimeError("missing image") + + monkeypatch.setattr( + local_automation.docker, "from_env", lambda timeout=8: FakeDockerEnv() + ) + monkeypatch.setattr( + local_automation.toml, + "load", + lambda _: { + "job": {"pool_name": "p1", "job_name": "j2"}, + "experiment": {"base_cmd": "python task.py", "exp_yaml": "grid.yaml"}, + }, + ) + + local_automation.run_experiment("exp.toml") + + assert fake.cont_name == "repo_image_tag" + assert fake.calls["add_tasks_from_yaml"][0]["file_path"] == "grid.yaml" + + +def test_local_run_tasks_client_creation_failure(monkeypatch): + monkeypatch.setattr( + local_automation, + "CloudClient", + lambda **kwargs: (_ for _ in ()).throw(RuntimeError("boom")), + ) + monkeypatch.setattr( + local_automation.toml, + "load", + lambda _: {"job": {"pool_name": "p", "job_name": "j"}, "task": []}, + ) + + assert local_automation.run_tasks("tasks.toml") is None + + +def test_local_run_tasks_missing_pool_name(monkeypatch): + fake = FakeLocalClient() + monkeypatch.setattr(local_automation, "CloudClient", lambda **kwargs: fake) + monkeypatch.setattr( + local_automation.toml, "load", lambda _: {"job": {}, "task": []} + ) + + assert local_automation.run_tasks("tasks.toml") is None + + +def test_local_run_tasks_no_container_reads_pool_and_monitor_false(monkeypatch): + fake = FakeLocalClient() + monkeypatch.setattr(local_automation, "CloudClient", lambda **kwargs: fake) + monkeypatch.setattr( + local_automation.Path, + "read_text", + lambda self: "{'image_name': 'repo/image:tag'}", + ) + monkeypatch.setattr( + local_automation.toml, + "load", + lambda _: { + "job": { + "pool_name": "p1", + "job_name": "j3", + "save_logs_to_blob": "logs", + "logs_folder": "folder", + "task_retries": 4, + "monitor_job": False, + }, + "upload": {"container_name": "cont", "files": ["x.txt"]}, + "task": [ + {"name": "a", "cmd": "echo a"}, + {"name": "b", "cmd": "echo b", "depends_on": ["a"]}, + ], + }, + ) + + assert local_automation.run_tasks("tasks.toml") is None + + assert fake.calls["create_job"][0]["task_retries"] == 4 + assert fake.calls["upload_files"][0]["location_in_blob"] == "" + assert fake.calls["add_task"][0]["run_dependent_tasks_on_fail"] is False + assert fake.calls["add_task"][1]["depends_on"] == ["tid-1"] + assert fake.calls["add_task"][1]["container_image_name"] == "repo_image_tag.j3" + assert fake.calls["monitor_job"] == [] + + +def test_batch_pool_service_init(monkeypatch, svc_mod): + attrs = { + "AZURE_TENANT_ID": "tenant", + "AZURE_SUBSCRIPTION_ID": "sub", + "AZURE_SP_CLIENT_ID": "cid", + "AZURE_CLIENT_SECRET": "secret", # pragma: allowlist secret + "AZURE_KEYVAULT_ENDPOINT": "kv", + "AZURE_KEYVAULT_SP_SECRET_ID": "sid", + "AZURE_RESOURCE_GROUP": "rg", + "AZURE_BLOB_STORAGE_ACCOUNT": "blob", + "AZURE_SUBNET_ID": "subnet", + "AZURE_BATCH_ACCOUNT": "batch", + } + monkeypatch.setattr(svc_mod, "dotenv_values", lambda _: attrs) + monkeypatch.setattr( + svc_mod.toml, "load", lambda _: {"Pool": {"parallel_pool_limit": "2"}} + ) + + cred = SimpleNamespace(compute_node_identity_reference="idref") + monkeypatch.setattr(svc_mod, "SPCredentialHandler", lambda **kwargs: cred) + monkeypatch.setattr(svc_mod, "get_batch_management_client", lambda c: "bmc") + + svc = svc_mod.CFABatchPoolService(".env", "job.toml") + assert svc.parallel_pool_limit == 2 + assert svc.batch_mgmt_client == "bmc" + + +def test_batch_pool_service_setup_pools_branches(monkeypatch, svc_mod): + svc = svc_mod.CFABatchPoolService.__new__(svc_mod.CFABatchPoolService) + svc.parallel_pool_limit = 2 + called = [] + + monkeypatch.setattr( + svc, + "_CFABatchPoolService__setup_pool", + lambda pool_name: called.append(pool_name), + ) + + svc.job_configuration = {"Pool": {"pool_name": "fixed"}} + svc.setup_pools() + assert called == ["fixed"] + + called.clear() + svc.job_configuration = {"Pool": {}} + svc.setup_pools(pools=["p1", "p2"]) + assert called == ["p1", "p2"] + + called.clear() + svc.job_configuration = {"Pool": {"pool_name_prefix": "pref_"}} + svc.setup_pools() + assert called == ["pref_0", "pref_1"] + + +def test_batch_pool_service_setup_pool_paths(monkeypatch, svc_mod): + svc = svc_mod.CFABatchPoolService.__new__(svc_mod.CFABatchPoolService) + svc.batch_pools = [] + svc.cred = SimpleNamespace( + azure_resource_group_name="rg", azure_batch_account="acct" + ) + svc.batch_mgmt_client = "bmc" + + create_calls = [] + monkeypatch.setattr(svc_mod.bh, "check_pool_exists", lambda *a, **k: True) + monkeypatch.setattr(svc, "_CFABatchPoolService__create_containers", lambda: "mc") + monkeypatch.setattr( + svc, "_CFABatchPoolService__create_pool_configuration", lambda n, m: "pc" + ) + monkeypatch.setattr( + svc, + "_CFABatchPoolService__create_pool", + lambda n, p: create_calls.append((n, p)), + ) + + svc._CFABatchPoolService__setup_pool("pool-a") + assert svc.batch_pools == ["pool-a"] + assert create_calls == [] + + monkeypatch.setattr(svc_mod.bh, "check_pool_exists", lambda *a, **k: False) + svc._CFABatchPoolService__setup_pool("pool-b") + assert ("pool-b", "pc") in create_calls + + +def test_batch_pool_service_create_containers(monkeypatch, svc_mod): + svc = svc_mod.CFABatchPoolService.__new__(svc_mod.CFABatchPoolService) + svc.job_configuration = {"Pool": {"input_mount": "in", "output_mount": "out"}} + svc.cred = SimpleNamespace( + azure_blob_storage_account="blobacct", + compute_node_identity_reference="idref", + ) + + seen = {} + + def fake_get_node_mount_config(**kwargs): + seen.update(kwargs) + return {"mount": "ok"} + + monkeypatch.setattr(svc_mod, "get_node_mount_config", fake_get_node_mount_config) + + out = svc._CFABatchPoolService__create_containers() + assert out == {"mount": "ok"} + assert seen["storage_containers"] == ["in", "out"] + + +def test_batch_pool_service_create_pool_configuration(monkeypatch, svc_mod): + svc = svc_mod.CFABatchPoolService.__new__(svc_mod.CFABatchPoolService) + svc.job_configuration = { + "Pool": { + "autoscale": "false", + "task_slots_per_node": "3", + "container_image_name": "repo/image:tag", + "vm_size": "Standard_D2s_v3", + } + } + svc.cred = SimpleNamespace( + azure_subnet_id="subnet", + azure_user_assigned_identity="uami", + azure_container_registry="registry", + ) + + fake_pool_cfg = SimpleNamespace( + deployment_configuration=SimpleNamespace( + virtual_machine_configuration=SimpleNamespace( + node_placement_configuration=None + ) + ) + ) + + monkeypatch.setattr( + svc_mod, "get_default_pool_config", lambda **kwargs: fake_pool_cfg + ) + monkeypatch.setattr( + svc, + "_CFABatchPoolService__setup_fixedscale_configuration", + lambda pool_config: pool_config, + ) + + assigned = [] + monkeypatch.setattr( + svc_mod, "assign_container_config", lambda p, c: assigned.append((p, c)) + ) + + monkeypatch.setattr( + svc_mod, + "models", + SimpleNamespace( + ContainerConfiguration=lambda **kwargs: SimpleNamespace(**kwargs), + NodePlacementConfiguration=lambda **kwargs: SimpleNamespace(**kwargs), + NodePlacementPolicyType=SimpleNamespace(regional="regional"), + ), + ) + + out = svc._CFABatchPoolService__create_pool_configuration("pool-a", {"mount": "ok"}) + assert out.task_slots_per_node == 3 + assert assigned + assert ( + out.deployment_configuration.virtual_machine_configuration.node_placement_configuration.policy + == "regional" + ) + + +def test_batch_pool_service_create_pool_and_delete(monkeypatch, svc_mod): + svc = svc_mod.CFABatchPoolService.__new__(svc_mod.CFABatchPoolService) + svc.cred = SimpleNamespace( + azure_resource_group_name="rg", azure_batch_account="acct" + ) + calls = [] + + class PoolAPI: + def create(self, **kwargs): + calls.append(kwargs) + + svc.batch_mgmt_client = SimpleNamespace(pool=PoolAPI()) + svc._CFABatchPoolService__create_pool("pool-ok", "cfg") + assert svc.pool_name == "pool-ok" + assert calls[0]["pool_name"] == "pool-ok" + + class FailingPoolAPI: + def create(self, **kwargs): + raise RuntimeError("bad") + + svc.batch_mgmt_client = SimpleNamespace(pool=FailingPoolAPI()) + with pytest.raises(RuntimeError, match="Failed to create pool"): + svc._CFABatchPoolService__create_pool("pool-fail", "cfg") + + svc.batch_pools = ["p1", "p2"] + deleted = [] + monkeypatch.setattr( + svc_mod.bh, "delete_pool", lambda **kwargs: deleted.append(kwargs["pool_name"]) + ) + assert svc.delete_all_pools() is True + assert deleted == ["p1", "p2"] + + +def test_batch_pool_service_step_parameters_with_explicit_pools(svc_mod): + svc = svc_mod.CFABatchPoolService.__new__(svc_mod.CFABatchPoolService) + svc.job_configuration = {"Job": {"docker_command": "python {task_input}"}} + svc.parallel_pool_limit = 4 + svc.batch_pools = ["pool0", "pool1"] + svc.attributes = {"k": "v"} + + params = svc.setup_step_parameters([1, 2, 3, 4], pools=["pool0", "pool1"]) + assert len(params) == 2 + assert params[0]["pool_name"] == "pool0" + assert params[1]["pool_name"] == "pool1" diff --git a/tests/test_local_modules_coverage.py b/tests/test_local_modules_coverage.py new file mode 100644 index 0000000..f283efa --- /dev/null +++ b/tests/test_local_modules_coverage.py @@ -0,0 +1,255 @@ +from pathlib import Path +from types import SimpleNamespace + +import pytest +import yaml + +from cfa.cloudops.local import _client as local_client +from cfa.cloudops.local import automation as local_automation +from cfa.cloudops.local import batch as local_batch +from cfa.cloudops.local import helpers as local_helpers + + +def test_local_batch_task_dependencies_and_repr(): + t1 = local_batch.Task("echo 1", id="t1") + t2 = local_batch.Task("echo 2", id="t2") + + t1.before(t2) + assert t1 in t2.deps + + t2.after(t1) + assert t1 in t2.deps + + t1.set_downstream(t2) + t2.set_upstream(t1) + + assert repr(t1) == "t1" + + +def test_local_batch_pool_and_job_models(): + pool = local_batch.Pool("pool", "image") + job = local_batch.Job("job", "pool", 1, True) + + assert pool.pool_id == "pool" + assert job.job_id == "job" + + +def test_local_helpers_add_job_and_create_container(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + + job = local_helpers.add_job("job1", "pool1") + assert job.job_id == "job1" + assert Path("tmp/jobs/job1.txt").exists() + + c = local_helpers.create_container("cont1") + assert c == "container_client" + assert Path("cont1").exists() + + +def test_local_helpers_upload_and_download_file(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + + Path("cont").mkdir() + src = Path("src.txt") + src.write_text("hello") + + local_helpers.upload_to_storage_container( + filepath=str(src), + location="folder", + container_name="cont", + verbose=False, + ) + assert Path("cont/folder/src.txt").read_text() == "hello" + + dest = Path("dest.txt") + local_helpers.download_file(None, "src.txt", str(dest), True, True) + assert dest.read_text() == "hello" + + +def test_local_helpers_upload_folder_filters(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + + Path("cont").mkdir() + Path("folder/sub").mkdir(parents=True) + Path("folder/a.txt").write_text("a") + Path("folder/sub/b.csv").write_text("b") + + uploaded = local_helpers.upload_folder( + folder="folder", + container_name="cont", + include_extensions=["txt"], + location_in_blob="blobdir", + ) + + assert any(x.endswith("a.txt") for x in uploaded) + assert Path("cont/blobdir/a.txt").exists() + + +def test_local_helpers_docker_helpers(monkeypatch): + class FakeImage: + def tag(self, name): + self.name = name + + fake_image = FakeImage() + + docker_env = SimpleNamespace( + ping=lambda: True, + images=SimpleNamespace( + get=lambda name: fake_image, + list=lambda: [SimpleNamespace(tags=["x:1"])], + ), + ) + + monkeypatch.setattr( + "cfa.cloudops.local.helpers.docker.from_env", lambda timeout=10: docker_env + ) + monkeypatch.setattr("cfa.cloudops.local.helpers.os.path.exists", lambda p: True) + monkeypatch.setattr("cfa.cloudops.local.helpers.sp.run", lambda *a, **k: None) + + out1 = local_helpers.package_and_upload_dockerfile("reg", "repo", "v1") + out2 = local_helpers.upload_docker_image("local:1", "reg", "repo", "v2") + + assert out1 == "reg.azurecr.io/repo:v1" + assert out2 == "reg.azurecr.io/repo:v2" + + +def test_local_helpers_yaml_and_walk(tmp_path, monkeypatch): + config = { + "param": [1, 2], + "flag(flag)": ["x", ""], + } + fpath = tmp_path / "args.yaml" + with open(fpath, "w") as f: + yaml.safe_dump(config, f) + + class FakeGrid: + def to_dicts(self): + return [{"param": 1, "flag(flag)": "x"}, {"param": 2, "flag(flag)": ""}] + + monkeypatch.setattr(local_helpers, "parse", lambda raw: FakeGrid()) + + args = local_helpers.get_args_from_yaml(str(fpath)) + cmds = local_helpers.get_tasks_from_yaml("python script.py", str(fpath)) + + assert len(args) > 0 + assert all(x.startswith("python script.py") for x in cmds) + + +def test_local_helpers_format_extensions_and_walk_folder(tmp_path): + assert local_helpers.format_extensions("txt") == [".txt"] + assert local_helpers.format_extensions([".txt", "csv"]) == [".txt", ".csv"] + + (tmp_path / "d").mkdir() + (tmp_path / "d" / "f.txt").write_text("x") + files = local_helpers.walk_folder(str(tmp_path / "d")) + assert any(x.endswith("f.txt") for x in files) + + +def test_local_cloudclient_basic_ops(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + + docker_env = SimpleNamespace( + ping=lambda: True, + images=SimpleNamespace( + get=lambda name: SimpleNamespace(short_id="img1"), + pull=lambda name: SimpleNamespace(tags=[name]), + ), + ) + + monkeypatch.setattr( + "cfa.cloudops.local._client.docker.from_env", lambda timeout=8: docker_env + ) + monkeypatch.setattr("cfa.cloudops.local._client.sp.run", lambda *a, **k: None) + + c = local_client.CloudClient() + + with pytest.raises(ValueError): + c.create_pool("p1", container_image_name=None) + + out = c.create_pool("pool1", container_image_name="python:3.11") + assert out["pool_id"] == "pool1" + assert Path("tmp/pools/pool1.txt").exists() + + c.create_job("job1", "pool1") + assert Path("tmp/jobs/job1.txt").exists() + + tid = c.add_task("job1", "echo hello") + assert isinstance(tid, int) + + c.create_blob_container("cont") + Path("cont/file.txt").write_text("x") + assert c.list_blob_files("cont") == ["file.txt"] + + +def test_local_cloudclient_blob_delete_and_yaml(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + + c = local_client.CloudClient() + Path("tmp/jobs").mkdir(parents=True) + Path("tmp/pools").mkdir(parents=True) + Path("tmp/jobs/jobx.txt").write_text("poolx 0 False") + Path("tmp/pools/poolx.txt").write_text("{'image_name': 'img:1', 'mount_str': ''}") + + monkeypatch.setattr("cfa.cloudops.local._client.sp.run", lambda *a, **k: None) + monkeypatch.setattr( + "cfa.cloudops.local._client.docker.from_env", + lambda: SimpleNamespace( + images=SimpleNamespace(get=lambda name: SimpleNamespace(short_id="i")) + ), + ) + + monkeypatch.setattr( + "cfa.cloudops.local._client.helpers.get_tasks_from_yaml", + lambda **k: ["echo 1", "echo 2"], + ) + monkeypatch.setattr(c, "add_task", lambda **k: 1) + + tasks = c.add_tasks_from_yaml("jobx", "python x.py", "cfg.yml") + assert tasks == [1, 1] + + Path("cont").mkdir() + Path("cont/a.txt").write_text("x") + c.delete_blob_file("a.txt", "cont") + assert not Path("cont/a.txt").exists() + + Path("cont/f").mkdir(parents=True) + Path("cont/f/b.txt").write_text("y") + c.delete_blob_folder("f", "cont") + assert not Path("cont/f").exists() + + +def test_local_automation_missing_pool_name(monkeypatch): + monkeypatch.setattr( + "cfa.cloudops.local.automation.toml.load", lambda _: {"job": {}} + ) + assert local_automation.run_experiment("exp.toml") is None + + +def test_local_automation_run_tasks_happy_path(monkeypatch): + config = { + "job": {"pool_name": "p1", "job_name": "j1"}, + "task": [ + {"name": "t1", "cmd": "echo 1"}, + {"name": "t2", "cmd": "echo 2", "depends_on": ["t1"]}, + ], + } + + fake_client = SimpleNamespace( + upload_folders=lambda **k: None, + upload_files=lambda **k: None, + create_job=lambda **k: None, + add_task=lambda **k: "task-id", + monitor_job=lambda *a, **k: None, + ) + + monkeypatch.setattr("cfa.cloudops.local.automation.toml.load", lambda _: config) + monkeypatch.setattr( + "cfa.cloudops.local.automation.CloudClient", + lambda dotenv_path=None: fake_client, + ) + monkeypatch.setattr( + "cfa.cloudops.local.automation.Path.read_text", + lambda self: "{'image_name': 'img:1'}", + ) + + assert local_automation.run_tasks("tasks.toml") is None diff --git a/tests/test_metaflow_decorator_more.py b/tests/test_metaflow_decorator_more.py new file mode 100644 index 0000000..e577272 --- /dev/null +++ b/tests/test_metaflow_decorator_more.py @@ -0,0 +1,301 @@ +import importlib +import string +import sys +from types import ModuleType, SimpleNamespace + +import pytest + + +@pytest.fixture +def deco_mod(monkeypatch): + # Stub imports required by cfa.cloudops.metaflow package initialization. + fake_examples = { + "examples": ModuleType("examples"), + "examples.metaflow": ModuleType("examples.metaflow"), + "examples.metaflow.azure_batch_decorator": ModuleType( + "examples.metaflow.azure_batch_decorator" + ), + "examples.metaflow.plugins": ModuleType("examples.metaflow.plugins"), + "examples.metaflow.plugins.metadata_providers": ModuleType( + "examples.metaflow.plugins.metadata_providers" + ), + "examples.metaflow.plugins.metadata_providers.local": ModuleType( + "examples.metaflow.plugins.metadata_providers.local" + ), + } + + class FakeAzureBatchDecorator: + pass + + class FakeLocalMetadataProvider: + pass + + fake_examples[ + "examples.metaflow.azure_batch_decorator" + ].AzureBatchDecorator = FakeAzureBatchDecorator + fake_examples[ + "examples.metaflow.plugins.metadata_providers.local" + ].LocalMetadataProvider = FakeLocalMetadataProvider + + for name, mod in fake_examples.items(): + monkeypatch.setitem(sys.modules, name, mod) + + # Stub StepDecorator base class for isolated decorator import. + metaflow_pkg = ModuleType("metaflow") + metaflow_decorators = ModuleType("metaflow.decorators") + + class DummyStepDecorator: + def __init__(self, *args, **kwargs): + pass + + metaflow_decorators.StepDecorator = DummyStepDecorator + monkeypatch.setitem(sys.modules, "metaflow", metaflow_pkg) + monkeypatch.setitem(sys.modules, "metaflow.decorators", metaflow_decorators) + + return importlib.import_module( + "cfa.cloudops.metaflow.custom_metaflow.plugins.decorators.cfa_azure_batch_decorator" + ) + + +def _attrs(): + return { + "AZURE_TENANT_ID": "tenant", + "AZURE_SUBSCRIPTION_ID": "sub", + "AZURE_SP_CLIENT_ID": "client", + "AZURE_CLIENT_SECRET": "secret", # pragma: allowlist secret + "AZURE_KEYVAULT_ENDPOINT": "kv", + "AZURE_KEYVAULT_SP_SECRET_ID": "sid", + "AZURE_RESOURCE_GROUP": "rg", + "AZURE_BATCH_ACCOUNT": "batch", + "AZURE_BLOB_STORAGE_ACCOUNT": "blob", + "AZURE_SUBNET_ID": "subnet", + "AZURE_USER_ASSIGNED_IDENTITY": "uami", + } + + +def test_generate_random_string_is_alnum(deco_mod): + out = deco_mod.generate_random_string(12) + assert len(out) == 12 + assert all(ch in (string.ascii_letters + string.digits) for ch in out) + + +def test_decorator_init_sets_clients(monkeypatch, deco_mod): + cred = SimpleNamespace() + monkeypatch.setattr(deco_mod, "SPCredentialHandler", lambda **kwargs: cred) + monkeypatch.setattr(deco_mod, "get_batch_service_client", lambda c: "batch-client") + monkeypatch.setattr( + deco_mod, "get_batch_management_client", lambda c: "batch-mgmt-client" + ) + + d = deco_mod.CFAAzureBatchDecorator( + pool_name="pool-a", + attributes=_attrs(), + job_configuration={"Job": {}, "Pool": {}}, + docker_command="python main.py", + task_parameters=[1, 2], + ) + + assert d.pool_name == "pool-a" + assert d.batch_client == "batch-client" + assert d.batch_mgmt_client == "batch-mgmt-client" + assert d.task_interval == 10 + assert d.docker_command == "python main.py" + assert d.task_parameters == [1, 2] + assert d.cred.azure_resource_group_name == "rg" + assert d.cred.azure_batch_account == "batch" + + +def test_private_create_job_builds_and_submits(monkeypatch, deco_mod): + d = deco_mod.CFAAzureBatchDecorator.__new__(deco_mod.CFAAzureBatchDecorator) + d.pool_name = "pool-a" + d.batch_client = "batch-client" + + class FakeConstraints: + def __init__(self, max_task_retry_count=None, max_wall_clock_time=None): + self.max_task_retry_count = max_task_retry_count + self.max_wall_clock_time = max_wall_clock_time + + class FakeJob: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + fake_models = SimpleNamespace( + BatchAllTasksCompleteMode=SimpleNamespace( + TERMINATE_JOB="term", NO_ACTION="none" + ), + BatchTaskFailureMode=SimpleNamespace(PERFORM_EXIT_OPTIONS_JOB_ACTION="exit"), + BatchJobConstraints=FakeConstraints, + BatchPoolInfo=lambda pool_id: SimpleNamespace(pool_id=pool_id), + BatchMetadataItem=lambda name, value: SimpleNamespace(name=name, value=value), + BatchJobCreateOptions=lambda **kwargs: FakeJob(**kwargs), + ) + monkeypatch.setattr(deco_mod, "batch_models", fake_models) + + submitted = {} + + def fake_create_job(client, job, exist_ok=False, verify_pool=True, verbose=False): + submitted.update( + { + "client": client, + "job": job, + "exist_ok": exist_ok, + "verify_pool": verify_pool, + "verbose": verbose, + } + ) + + monkeypatch.setattr(deco_mod, "create_job", fake_create_job) + + d._CFAAzureBatchDecorator__create_job( + job_name="my job", + task_retries=2, + mark_complete_after_tasks_run=True, + timeout=5, + uses_deps=False, + exist_ok=True, + verify_pool=False, + verbose=True, + ) + + job = submitted["job"] + assert submitted["client"] == "batch-client" + assert submitted["exist_ok"] is True + assert submitted["verify_pool"] is False + assert submitted["verbose"] is True + assert job.id == "myjob" + assert job.pool_info.pool_id == "pool-a" + assert job.uses_task_dependencies is False + assert job.all_tasks_complete_mode == "term" + assert job.constraints.max_task_retry_count == 2 + assert job.constraints.max_wall_clock_time is not None + + +def test_add_task_delegates_and_overrides_logs_folder(monkeypatch, deco_mod): + d = deco_mod.CFAAzureBatchDecorator.__new__(deco_mod.CFAAzureBatchDecorator) + d.pool_name = "pool-a" + d.cred = SimpleNamespace( + azure_resource_group_name="rg", azure_batch_account="batch" + ) + d.batch_mgmt_client = "batch-mgmt" + d.batch_client = "batch-client" + + monkeypatch.setattr( + deco_mod.batch_helpers, "get_pool_mounts", lambda *a, **k: ["/mnt/input"] + ) + seen = {} + + def fake_add_task(**kwargs): + seen.update(kwargs) + return "tid-1" + + monkeypatch.setattr(deco_mod.batch_helpers, "add_task", fake_add_task) + + tid = d.add_task( + job_name="job-1", + command_line="echo hi", + save_logs_to_blob="logs", + logs_folder="custom", + depends_on=["a"], + run_dependent_tasks_on_fail=True, + container_image_name="img:1", + timeout=9, + ) + + assert tid == "tid-1" + assert seen["job_name"] == "job-1" + assert seen["logs_folder"] == "stdout_stderr" + assert seen["mounts"] == ["/mnt/input"] + assert seen["depends_on"] == ["a"] + assert seen["batch_client"] == "batch-client" + assert seen["full_container_name"] == "img:1" + + +def test_fetch_or_create_job_reuse_and_create(monkeypatch, deco_mod): + d = deco_mod.CFAAzureBatchDecorator.__new__(deco_mod.CFAAzureBatchDecorator) + d.job_configuration = {"Job": {"job_id": "jid", "job_id_prefix": "pref-"}} + d.batch_client = "batch-client" + + monkeypatch.setattr(deco_mod, "generate_random_string", lambda length: "abcde") + + monkeypatch.setattr( + deco_mod.batch_helpers, "check_job_exists", lambda job_id, client: True + ) + assert d.fetch_or_create_job() == "pref-abcde" + + calls = [] + monkeypatch.setattr( + deco_mod.batch_helpers, "check_job_exists", lambda job_id, client: False + ) + monkeypatch.setattr( + d, + "_CFAAzureBatchDecorator__create_job", + lambda **kwargs: calls.append(kwargs), + ) + assert d.fetch_or_create_job() == "pref-abcde" + assert calls[0]["job_name"] == "pref-abcde" + assert calls[0]["mark_complete_after_tasks_run"] is True + + +def test_wrapper_submits_tasks_and_calls_function(monkeypatch, deco_mod): + d = deco_mod.CFAAzureBatchDecorator.__new__(deco_mod.CFAAzureBatchDecorator) + d.job_configuration = { + "Job": {"parent_task": "p1,p2"}, + "Pool": {"container_image_name": "img:tag"}, + } + d.task_interval = 0 + d.docker_command = "python run.py --input {task_input} --jid {job_id}" + d.task_parameters = ["a", "b"] + + monkeypatch.setattr(d, "fetch_or_create_job", lambda: "jobx") + monkeypatch.setattr(deco_mod.time, "sleep", lambda _: None) + monkeypatch.setattr(deco_mod, "generate_random_string", lambda length: "XYZ") + + calls = [] + + def fake_add_task(**kwargs): + calls.append(kwargs) + return f"tid-{len(calls)}" + + monkeypatch.setattr(d, "add_task", fake_add_task) + + @d + def run(x, y): + return x + y + + result = run(1, 2) + assert result == 3 + assert len(calls) == 2 + assert calls[0]["job_name"] == "jobx" + assert calls[0]["depends_on"] == ["p1", "p2"] + assert calls[0]["container_image_name"] == "img:tag" + assert "--input a" in calls[0]["command_line"] + assert calls[0]["name_suffix"] == "jobx_task_XYZ_" + assert d.task_id == "tid-2" + + +def test_wrapper_uses_default_container_when_not_configured(monkeypatch, deco_mod): + d = deco_mod.CFAAzureBatchDecorator.__new__(deco_mod.CFAAzureBatchDecorator) + d.job_configuration = {"Job": {}, "Pool": {}} + d.task_interval = 0 + d.docker_command = "echo {task_input} {job_id}" + d.task_parameters = [1] + + monkeypatch.setattr(d, "fetch_or_create_job", lambda: "joby") + monkeypatch.setattr(deco_mod.time, "sleep", lambda _: None) + monkeypatch.setattr(deco_mod, "generate_random_string", lambda length: "QWE") + + seen = {} + + def fake_add_task(**kwargs): + seen.update(kwargs) + return "tid-1" + + monkeypatch.setattr(d, "add_task", fake_add_task) + + @d + def run(): + return "ok" + + assert run() == "ok" + assert seen["depends_on"] is None + assert seen["container_image_name"] == "python:latest" diff --git a/tests/test_scripts.py b/tests/test_scripts.py index a53f266..254b4bd 100644 --- a/tests/test_scripts.py +++ b/tests/test_scripts.py @@ -1,3 +1,4 @@ +import pytest from shared_fixtures import FAKE_COMMANDLINE import cfa.cloudops.scripts as scripts @@ -80,3 +81,240 @@ def test_download_after_job(mocker, monkeypatch): "cfa.cloudops._cloudclient.CloudClient.download_after_job", return_value=None ) scripts.download_after_job() + + +def test_hello(monkeypatch, capsys): + monkeypatch.setattr("sys.argv", ["script_name.py", "--name", "Tester"]) + scripts.hello() + captured = capsys.readouterr() + assert "Hello, Tester!" in captured.out + + +def test_create_blob_container(mocker, monkeypatch): + monkeypatch.setattr( + "sys.argv", FAKE_COMMANDLINE + ["--container_name", "my-container"] + ) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.__init__", return_value=None) + mocker.patch( + "cfa.cloudops._cloudclient.CloudClient.create_blob_container", + return_value=None, + ) + scripts.create_blob_container() + + +def test_monitor_job(mocker, monkeypatch): + monkeypatch.setattr( + "sys.argv", FAKE_COMMANDLINE + ["--job_name", "job-1", "--download_job_stats"] + ) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.__init__", return_value=None) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.monitor_job", return_value=None) + scripts.monitor_job() + + +def test_check_job_status(mocker, monkeypatch, capsys): + monkeypatch.setattr("sys.argv", FAKE_COMMANDLINE + ["--job_name", "job-1"]) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.__init__", return_value=None) + mocker.patch( + "cfa.cloudops._cloudclient.CloudClient.check_job_status", + return_value="completed", + ) + scripts.check_job_status() + captured = capsys.readouterr() + assert "completed" in captured.out + + +def test_delete_job(mocker, monkeypatch): + monkeypatch.setattr("sys.argv", FAKE_COMMANDLINE + ["--job_name", "job-1"]) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.__init__", return_value=None) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.delete_job", return_value=None) + scripts.delete_job() + + +def test_package_and_upload_dockerfile(mocker, monkeypatch): + monkeypatch.setattr( + "sys.argv", + FAKE_COMMANDLINE + + [ + "--registry_name", + "reg", + "--repo_name", + "repo", + "--tag", + "v1", + "--path_to_dockerfile", + "./Dockerfile", + "--use_device_code", + ], + ) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.__init__", return_value=None) + mocker.patch( + "cfa.cloudops._cloudclient.CloudClient.package_and_upload_dockerfile", + return_value=None, + ) + scripts.package_and_upload_dockerfile() + + +def test_upload_docker_image(mocker, monkeypatch): + monkeypatch.setattr( + "sys.argv", + FAKE_COMMANDLINE + + [ + "--image_name", + "local:latest", + "--registry_name", + "reg", + "--repo_name", + "repo", + "--tag", + "v2", + "--use_device_code", + ], + ) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.__init__", return_value=None) + mocker.patch( + "cfa.cloudops._cloudclient.CloudClient.upload_docker_image", + return_value=None, + ) + scripts.upload_docker_image() + + +def test_download_file(mocker, monkeypatch): + monkeypatch.setattr( + "sys.argv", + FAKE_COMMANDLINE + + [ + "--container_name", + "my-container", + "--blob_name", + "path/file.txt", + "--destination_path", + "./file.txt", + "--check_size", + ], + ) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.__init__", return_value=None) + mocker.patch( + "cfa.cloudops._cloudclient.CloudClient.download_file", return_value=None + ) + scripts.download_file() + + +def test_download_folder(mocker, monkeypatch): + monkeypatch.setattr( + "sys.argv", + FAKE_COMMANDLINE + + [ + "--src_path", + "my-src-path", + "--dest_path", + "./downloads", + "--container_name", + "my-container", + "--include_extensions", + ".txt", + "--check_size", + ], + ) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.__init__", return_value=None) + mocker.patch( + "cfa.cloudops._cloudclient.CloudClient.download_folder", return_value=None + ) + scripts.download_folder() + + +def test_delete_pool(mocker, monkeypatch): + monkeypatch.setattr("sys.argv", FAKE_COMMANDLINE + ["--pool_name", "pool-1"]) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.__init__", return_value=None) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.delete_pool", return_value=None) + scripts.delete_pool() + + +def test_list_blob_files(mocker, monkeypatch, capsys): + monkeypatch.setattr( + "sys.argv", FAKE_COMMANDLINE + ["--container_name", "my-container"] + ) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.__init__", return_value=None) + mocker.patch( + "cfa.cloudops._cloudclient.CloudClient.list_blob_files", + return_value=["a.txt", "b.txt"], + ) + scripts.list_blob_files() + captured = capsys.readouterr() + assert "a.txt" in captured.out + assert "b.txt" in captured.out + + +def test_delete_blob_file(mocker, monkeypatch): + monkeypatch.setattr( + "sys.argv", + FAKE_COMMANDLINE + + ["--container_name", "my-container", "--blob_name", "a/file.txt"], + ) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.__init__", return_value=None) + mocker.patch( + "cfa.cloudops._cloudclient.CloudClient.delete_blob_file", return_value=None + ) + scripts.delete_blob_file() + + +def test_delete_blob_folder(mocker, monkeypatch): + monkeypatch.setattr( + "sys.argv", + FAKE_COMMANDLINE + + ["--container_name", "my-container", "--blob_folder_name", "folder/a"], + ) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.__init__", return_value=None) + mocker.patch( + "cfa.cloudops._cloudclient.CloudClient.delete_blob_folder", return_value=None + ) + scripts.delete_blob_folder() + + +def test_download_job_stats(mocker, monkeypatch): + monkeypatch.setattr( + "sys.argv", + FAKE_COMMANDLINE + ["--job_name", "job-1", "--file_name", "stats.csv"], + ) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.__init__", return_value=None) + mocker.patch( + "cfa.cloudops._cloudclient.CloudClient.download_job_stats", return_value=None + ) + scripts.download_job_stats() + + +def test_add_tasks_from_yaml(mocker, monkeypatch): + monkeypatch.setattr( + "sys.argv", + FAKE_COMMANDLINE + + [ + "--job_name", + "job-1", + "--base_cmd", + "python main.py", + "--file_path", + "tasks.yaml", + ], + ) + mocker.patch("cfa.cloudops._cloudclient.CloudClient.__init__", return_value=None) + mocker.patch( + "cfa.cloudops._cloudclient.CloudClient.add_tasks_from_yaml", return_value=None + ) + scripts.add_tasks_from_yaml() + + +def test_generate_sample_env(monkeypatch, tmp_path, capsys): + monkeypatch.chdir(tmp_path) + scripts.generate_sample_env() + generated = tmp_path / "cloudops-sample.env" + assert generated.exists() + assert "AZURE_BATCH_ACCOUNT" in generated.read_text() + captured = capsys.readouterr() + assert "created successfully" in captured.out + + +def test_test_entrypoint(mocker, monkeypatch): + monkeypatch.setattr("sys.argv", ["script_name.py", "-q"]) + mocker.patch("pytest.main", return_value=0) + with pytest.raises(SystemExit) as exc: + scripts.test() + assert exc.value.code == 0 diff --git a/tests/test_zero_modules_more.py b/tests/test_zero_modules_more.py new file mode 100644 index 0000000..75fd070 --- /dev/null +++ b/tests/test_zero_modules_more.py @@ -0,0 +1,287 @@ +import importlib +import sys +from types import ModuleType, SimpleNamespace + +import pytest + +from cfa.cloudops import _containerappclient as container_mod +from cfa.cloudops import _function_app_client as func_mod +from cfa.cloudops import autoscale + + +def test_autoscale_formulas_exposed(): + assert "maxNumberofVMs = 10" in autoscale.dev_autoscale_formula + assert "maxNumberofVMs = 25" in autoscale.prod_autoscale_formula + assert ( + "$NodeDeallocationOption = taskcompletion;" in autoscale.dev_autoscale_formula + ) + + +def test_containerappclient_methods_without_constructor(): + client = container_mod.ContainerAppClient.__new__(container_mod.ContainerAppClient) + client.resource_group = "rg" + client.job_name = "job1" + + c2 = SimpleNamespace(name="job2") + job_template = SimpleNamespace( + containers=[ + SimpleNamespace( + name="job1", + image="img:1", + command=["python"], + args=["x.py"], + env=[{"A": "B"}], + resources={"cpu": 1}, + ) + ] + ) + job_info = SimpleNamespace(name="job1", template=job_template) + job_info.as_dict = lambda: {"name": "job1"} + + jobs = SimpleNamespace( + list_by_resource_group=lambda rg: [job_info, c2], + begin_start=lambda **kwargs: SimpleNamespace(), + begin_stop_execution=lambda **kwargs: SimpleNamespace(result=lambda: "ok"), + ) + client.client = SimpleNamespace(jobs=jobs) + + assert client.list_jobs() == ["job1", "job2"] + assert client.check_job_exists("job1") is True + assert client.check_job_exists("missing") is False + assert client.get_job_info() == {"name": "job1"} + + info = client.get_command_info("job1") + assert info[0]["image"] == "img:1" + + client.start_job(job_name="job1") + + +@pytest.mark.parametrize( + "kwargs, expected", + [ + ({"command": "echo"}, "Command must be in list format."), + ({"args": "arg"}, "Args must be in list format."), + ({"env": ["A=B"]}, "Env must be in dict format."), + ], +) +def test_containerappclient_start_job_validation_errors(kwargs, expected): + client = container_mod.ContainerAppClient.__new__(container_mod.ContainerAppClient) + client.resource_group = "rg" + client.job_name = "job1" + client.client = SimpleNamespace( + jobs=SimpleNamespace( + list_by_resource_group=lambda rg: [], begin_start=lambda **k: None + ) + ) + + with pytest.raises(ValueError, match=expected): + client.start_job(**kwargs) + + +def test_containerappclient_start_job_with_overrides(monkeypatch): + env_objs = [] + + class FakeEnv: + def __init__(self, name=None, value=None, secret_ref=None): + env_objs.append((name, value, secret_ref)) + + monkeypatch.setattr(container_mod, "EnvironmentVar", FakeEnv) + monkeypatch.setattr( + container_mod, + "JobExecutionContainer", + lambda **k: SimpleNamespace(**k), + ) + monkeypatch.setattr( + container_mod, + "JobExecutionTemplate", + lambda **k: SimpleNamespace(**k), + ) + + client = container_mod.ContainerAppClient.__new__(container_mod.ContainerAppClient) + client.resource_group = "rg" + client.job_name = "job1" + + j = SimpleNamespace( + name="job1", + template=SimpleNamespace( + containers=[SimpleNamespace(name="c1", image="img", resources={"cpu": 1})] + ), + ) + started = {} + + def begin_start(**kwargs): + started.update(kwargs) + return SimpleNamespace() + + client.client = SimpleNamespace( + jobs=SimpleNamespace( + list_by_resource_group=lambda rg: [j], + begin_start=begin_start, + ) + ) + + client.start_job( + command=["python"], + args=["main.py"], + env={"A": "B"}, + secret_ref={"S": "secret"}, + ) + + assert ("A", "B", None) in env_objs + assert ("S", None, "secret") in env_objs + assert started["job_name"] == "job1" + + +def test_containerappclient_stop_job_error_path(): + client = container_mod.ContainerAppClient.__new__(container_mod.ContainerAppClient) + client.resource_group = "rg" + + client.client = SimpleNamespace( + jobs=SimpleNamespace( + begin_stop_execution=lambda **k: (_ for _ in ()).throw(RuntimeError("boom")) + ) + ) + + assert client.stop_job("j", "e") is None + + +def test_function_app_classmethods_and_validation(monkeypatch): + fake_cfg = SimpleNamespace( + additional_properties={"tags": ["one"]}, + health_check_path="/health", + ) + monkeypatch.setattr( + func_mod.FunctionAppClient, "get_configuration", lambda *a, **k: fake_cfg + ) + + assert func_mod.FunctionAppClient.get_tags("f") == ["one"] + assert func_mod.FunctionAppClient.get_health_check_flag("f") is True + + monkeypatch.delenv("AZURE_RESOURCE_GROUP", raising=False) + monkeypatch.delenv("AZURE_SUBSCRIPTION_ID", raising=False) + with pytest.raises(ValueError, match="Resource group"): + func_mod.FunctionAppClient.list_functions( + "f", resource_group=None, subscription_id="sub" + ) + + monkeypatch.setenv("AZURE_RESOURCE_GROUP", "rg") + monkeypatch.delenv("AZURE_SUBSCRIPTION_ID", raising=False) + with pytest.raises(ValueError, match="Subscription ID"): + func_mod.FunctionAppClient.list_slots( + "f", resource_group=None, subscription_id=None + ) + + +def test_function_app_init_and_database_connection(monkeypatch): + monkeypatch.setattr(func_mod, "EnvCredentialHandler", lambda **k: "envcred") + monkeypatch.setattr(func_mod, "DefaultCredentialHandler", lambda **k: "defaultcred") + monkeypatch.setattr(func_mod, "SPCredentialHandler", lambda **k: "spcred") + + c_env = func_mod.FunctionAppClient(function_app_name="f") + c_def = func_mod.FunctionAppClient(function_app_name="f", use_federated=True) + c_sp = func_mod.FunctionAppClient(function_app_name="f", use_sp=True) + + assert c_env.method == "env" + assert c_def.method == "default" + assert c_sp.method == "sp" + + sql_calls = [] + + class FakeConn: + def sql(self, q): + sql_calls.append(q) + return self + + monkeypatch.setattr( + func_mod.duckdb, "connect", lambda database=":memory:": FakeConn() + ) + + c_env.cred = SimpleNamespace( + azure_tenant_id="tenant", + azure_subscription_id="sub", + azure_client_id="cid", + azure_client_secret="secret", # pragma: allowlist secret + azure_blob_storage_account="storageacct", + ) + conn = c_env._get_database_connection() + assert conn is not None + assert any("CREATE SECRET" in q for q in sql_calls) + + +def test_metaflow_imports_and_decorator_basics(monkeypatch): + fake_examples = { + "examples": ModuleType("examples"), + "examples.metaflow": ModuleType("examples.metaflow"), + "examples.metaflow.azure_batch_decorator": ModuleType( + "examples.metaflow.azure_batch_decorator" + ), + "examples.metaflow.plugins": ModuleType("examples.metaflow.plugins"), + "examples.metaflow.plugins.metadata_providers": ModuleType( + "examples.metaflow.plugins.metadata_providers" + ), + "examples.metaflow.plugins.metadata_providers.local": ModuleType( + "examples.metaflow.plugins.metadata_providers.local" + ), + } + + class FakeAzureBatchDecorator: + pass + + class FakeLocalMetadataProvider: + pass + + fake_examples[ + "examples.metaflow.azure_batch_decorator" + ].AzureBatchDecorator = FakeAzureBatchDecorator + fake_examples[ + "examples.metaflow.plugins.metadata_providers.local" + ].LocalMetadataProvider = FakeLocalMetadataProvider + + for name, mod in fake_examples.items(): + monkeypatch.setitem(sys.modules, name, mod) + + metaflow_mod = importlib.reload(importlib.import_module("cfa.cloudops.metaflow")) + assert FakeAzureBatchDecorator in metaflow_mod.STEP_DECORATORS + assert FakeLocalMetadataProvider in metaflow_mod.METADATA_PROVIDERS + + # Stub minimal metaflow.decorators for decorator module import. + metaflow_pkg = ModuleType("metaflow") + metaflow_decorators = ModuleType("metaflow.decorators") + + class DummyStepDecorator: + def __init__(self, *args, **kwargs): + pass + + metaflow_decorators.StepDecorator = DummyStepDecorator + monkeypatch.setitem(sys.modules, "metaflow", metaflow_pkg) + monkeypatch.setitem(sys.modules, "metaflow.decorators", metaflow_decorators) + + deco_mod = importlib.reload( + importlib.import_module( + "cfa.cloudops.metaflow.custom_metaflow.plugins.decorators.cfa_azure_batch_decorator" + ) + ) + + assert len(deco_mod.generate_random_string(6)) == 6 + + +def test_cfa_batch_pool_service_setup_step_parameters(monkeypatch): + svc_mod = importlib.import_module( + "cfa.cloudops.metaflow.custom_metaflow.cfa_batch_pool_service" + ) + svc = svc_mod.CFABatchPoolService.__new__(svc_mod.CFABatchPoolService) + svc.job_configuration = {"Job": {"docker_command": "python {task_input}"}} + svc.parallel_pool_limit = 2 + svc.batch_pools = ["pool0", "pool1"] + svc.attributes = {"AZURE_SUBSCRIPTION_ID": "sub"} + + params = svc.setup_step_parameters([1, 2, 3, 4]) + assert len(params) == 2 + assert params[0]["pool_name"] == "pool0" + + +def test_decorators_init_exports(): + mod = importlib.import_module( + "cfa.cloudops.metaflow.custom_metaflow.plugins.decorators" + ) + assert "cfa_azure_batch" in mod.decorators