|
| 1 | +import re |
| 2 | +from unittest import mock |
| 3 | + |
| 4 | +import pytest |
| 5 | + |
| 6 | +from parsl.providers.kubernetes.kube import KubernetesProvider |
| 7 | + |
| 8 | +_MOCK_BASE = "parsl.providers.kubernetes.kube" |
| 9 | + |
| 10 | + |
| 11 | +@pytest.fixture(autouse=True) |
| 12 | +def mock_kube_config(): |
| 13 | + with mock.patch(f"{_MOCK_BASE}.config") as mock_config: |
| 14 | + mock_config.load_kube_config.return_value = None |
| 15 | + yield mock_config |
| 16 | + |
| 17 | + |
| 18 | +@pytest.fixture |
| 19 | +def mock_kube_client(): |
| 20 | + mock_client = mock.MagicMock() |
| 21 | + with mock.patch(f"{_MOCK_BASE}.client.CoreV1Api") as mock_api: |
| 22 | + mock_api.return_value = mock_client |
| 23 | + yield mock_client |
| 24 | + |
| 25 | + |
| 26 | +@pytest.mark.local |
| 27 | +def test_submit_happy_path(mock_kube_client: mock.MagicMock): |
| 28 | + image = "test-image" |
| 29 | + namespace = "test-namespace" |
| 30 | + cmd_string = "test-command" |
| 31 | + volumes = [("test-volume", "test-mount-path")] |
| 32 | + service_account_name = "test-service-account" |
| 33 | + annotations = {"test-annotation": "test-value"} |
| 34 | + max_cpu = 2 |
| 35 | + max_mem = "2Gi" |
| 36 | + init_cpu = 1 |
| 37 | + init_mem = "1Gi" |
| 38 | + provider = KubernetesProvider( |
| 39 | + image=image, |
| 40 | + persistent_volumes=volumes, |
| 41 | + namespace=namespace, |
| 42 | + service_account_name=service_account_name, |
| 43 | + annotations=annotations, |
| 44 | + max_cpu=max_cpu, |
| 45 | + max_mem=max_mem, |
| 46 | + init_cpu=init_cpu, |
| 47 | + init_mem=init_mem, |
| 48 | + ) |
| 49 | + |
| 50 | + job_name = "test.job.name" |
| 51 | + job_id = provider.submit(cmd_string=cmd_string, tasks_per_node=1, job_name=job_name) |
| 52 | + |
| 53 | + assert job_id in provider.resources |
| 54 | + assert mock_kube_client.create_namespaced_pod.call_count == 1 |
| 55 | + |
| 56 | + call_args = mock_kube_client.create_namespaced_pod.call_args[1] |
| 57 | + pod = call_args["body"] |
| 58 | + container = pod.spec.containers[0] |
| 59 | + volume = container.volume_mounts[0] |
| 60 | + |
| 61 | + assert image == container.image |
| 62 | + assert namespace == call_args["namespace"] |
| 63 | + assert any(cmd_string in arg for arg in container.args) |
| 64 | + assert volumes[0] == (volume.name, volume.mount_path) |
| 65 | + assert service_account_name == pod.spec.service_account_name |
| 66 | + assert annotations == pod.metadata.annotations |
| 67 | + assert str(max_cpu) == container.resources.limits["cpu"] |
| 68 | + assert max_mem == container.resources.limits["memory"] |
| 69 | + assert str(init_cpu) == container.resources.requests["cpu"] |
| 70 | + assert init_mem == container.resources.requests["memory"] |
| 71 | + assert job_id == pod.metadata.labels["job_id"] |
| 72 | + assert job_id == container.name |
| 73 | + assert f"{job_name}.{job_id}" == pod.metadata.name |
| 74 | + |
| 75 | + |
| 76 | +@pytest.mark.local |
| 77 | +@mock.patch(f"{_MOCK_BASE}.KubernetesProvider._create_pod") |
| 78 | +def test_submit_pod_name_includes_job_id(mock_create_pod: mock.MagicMock): |
| 79 | + provider = KubernetesProvider(image="test-image") |
| 80 | + |
| 81 | + job_name = "a." * 122 + "a" * 9 # Max length for pod name (253 chars) |
| 82 | + job_id = provider.submit(cmd_string="test-command", tasks_per_node=1, job_name=job_name) |
| 83 | + |
| 84 | + expected_pod_name = job_name[:-len(job_id) - 1] + job_id |
| 85 | + actual_pod_name = mock_create_pod.call_args[1]["pod_name"] |
| 86 | + assert not re.search(r'\.{2,}', actual_pod_name), "Pod name should not have consecutive dots" |
| 87 | + assert expected_pod_name == actual_pod_name |
| 88 | + |
| 89 | + |
| 90 | +@pytest.mark.local |
| 91 | +@mock.patch(f"{_MOCK_BASE}.KubernetesProvider._create_pod") |
| 92 | +def test_submit_empty_job_name(mock_create_pod: mock.MagicMock): |
| 93 | + provider = KubernetesProvider(image="test-image") |
| 94 | + job_id = provider.submit(cmd_string="test-command", tasks_per_node=1, job_name="") |
| 95 | + assert job_id == mock_create_pod.call_args[1]["pod_name"] |
0 commit comments