def test_kubernetes_cancel_after_ttl(
k8s_dispatcher: KubernetesDispatcher,
benchmark_event_dummy_payload: BenchmarkEvent,
k8s_test_client: KubernetesTestUtilsClient,
fetcher_job_config: FetcherJobConfig,
):
# No dispatch - should not throw
k8s_dispatcher.cancel_all(benchmark_event_dummy_payload.client_id, benchmark_event_dummy_payload.action_id)
def test_cleanup(
k8s_dispatcher: KubernetesDispatcher,
mock_batch_api_instance: kubernetes.client.BatchV1Api,
mock_core_api_instance: kubernetes.client.CoreV1Api,
mock_k8s_config,
):
k8s_dispatcher.cleanup(DATA_SET, BENCHMARK_EVENT)
_verify_k8s_all_delete(mock_batch_api_instance, mock_core_api_instance)
def test_cancel_all_actions(
k8s_dispatcher: KubernetesDispatcher,
mock_batch_api_instance: kubernetes.client.BatchV1Api,
mock_core_api_instance: kubernetes.client.CoreV1Api,
mock_k8s_config,
):
k8s_dispatcher.cancel_all(CLIENT_ID)
_verify_k8s_all_delete(mock_batch_api_instance, mock_core_api_instance)
def test_kubernetes_client(
k8s_dispatcher: KubernetesDispatcher,
benchmark_event_dummy_payload: BenchmarkEvent,
k8s_test_client: KubernetesTestUtilsClient,
fetcher_job_config: FetcherJobConfig,
size_info: ContentSizeInfo,
):
data_set = DownloadableContent(src=SOMEDATA_BIG, path="/mount/path", dst=S3_DST, md5=None, size_info=size_info)
k8s_dispatcher.dispatch_fetch(data_set, benchmark_event_dummy_payload, "/data/sets/fake")
_wait_for_k8s_objects_exist(benchmark_event_dummy_payload, fetcher_job_config, k8s_test_client, size_info)
def k8s_dispatcher(fetcher_service_config: FetcherServiceConfig):
return KubernetesDispatcher(
SERVICE_NAME,
fetcher_service_config.kubeconfig,
fetcher_service_config.zookeeper_ensemble_hosts,
fetcher_service_config.fetcher_job,
)
def k8s_dispatcher(
mock_k8s_config,
mock_batch_api_instance: kubernetes.client.BatchV1Api,
mock_core_api_instance: kubernetes.client.CoreV1Api,
):
return KubernetesDispatcher(
SERVICE_NAME, zk_ensemble=ZOOKEEPER_ENSEMBLE_HOSTS, kubeconfig=None, fetcher_job=FETCHER_JOB_CONFIG
)
def test_call_dispatcher(
k8s_dispatcher: KubernetesDispatcher,
mock_batch_api_instance: kubernetes.client.BatchV1Api,
mock_core_api_instance: kubernetes.client.CoreV1Api,
mock_k8s_config,
data_set: DownloadableContent,
size_info: ContentSizeInfo,
):
data_set_with_size = dataclasses.replace(data_set, size_info=size_info)
k8s_dispatcher.dispatch_fetch(data_set_with_size, BENCHMARK_EVENT, ZK_NODE_PATH)
mock_batch_api_instance.create_namespaced_job.assert_called_once()
job_args, _ = mock_batch_api_instance.create_namespaced_job.call_args
namespace, job = job_args
validate_namespaced_job(namespace, job, data_set)
def create_download_manager(zookeeper_ensemble_hosts: str, kubeconfig: str,
fetcher_job: FetcherJobConfig):
zk_client = KazooClient(zookeeper_ensemble_hosts)
job_dispatcher = KubernetesDispatcher(SERVICE_NAME, kubeconfig,
zookeeper_ensemble_hosts,
fetcher_job)
lock_manager = DistributedRWLockManager(zk_client, LOCK_MANAGER_PREFIX,
get_lock_name)
return DownloadManager(zk_client, job_dispatcher, lock_manager)
def test_get_label_selector():
assert (
KubernetesDispatcher.get_label_selector(SERVICE_NAME, CLIENT_ID, ACTION_ID, DATA_SET)
== ServiceLabels.get_label_selector(SERVICE_NAME, CLIENT_ID, ACTION_ID)
+ f",{KubernetesDispatcher.DATA_SET_HASH_LABEL}={md5sum(DATA_SET.src)}"
)
def test_call_dispatcher_no_size_infodef(k8s_dispatcher: KubernetesDispatcher):
with pytest.raises(ValueError):
k8s_dispatcher.dispatch_fetch(DATA_SET, BENCHMARK_EVENT, ZK_NODE_PATH)
def test_kubernetes_init_standalone(mock_k8s_client, mock_k8s_config):
KubernetesDispatcher(
SERVICE_NAME, zk_ensemble=ZOOKEEPER_ENSEMBLE_HOSTS, kubeconfig=KUBECONFIG, fetcher_job=FETCHER_JOB_CONFIG
)
mock_k8s_config.load_kube_config.assert_called_with(KUBECONFIG)
validate_client_calls(mock_k8s_client)
def test_kubernetes_init_in_cluster(mock_k8s_client, mock_k8s_config):
KubernetesDispatcher(
SERVICE_NAME, zk_ensemble=ZOOKEEPER_ENSEMBLE_HOSTS, kubeconfig=None, fetcher_job=FETCHER_JOB_CONFIG
)
mock_k8s_config.load_incluster_config.assert_called_once()
validate_client_calls(mock_k8s_client)