def test_cancel_no_node(kubernetes_job_starter: DownloadDispatcher,
zoo_keeper_client_no_node: KazooClient,
mock_lock_manager: RWLockManager):
download_mgr = DownloadManager(zoo_keeper_client_no_node,
kubernetes_job_starter, mock_lock_manager,
data_set_to_path)
download_mgr.cancel(CLIENT_ID, ACTION_ID)
zoo_keeper_client_no_node.set.assert_not_called()
def test_cancel_conflict(
kubernetes_job_starter: DownloadDispatcher,
zoo_keeper_client_with_conflict: KazooClient,
mock_lock_manager: RWLockManager,
):
download_mgr = DownloadManager(zoo_keeper_client_with_conflict,
kubernetes_job_starter, mock_lock_manager,
data_set_to_path)
download_mgr.cancel(CLIENT_ID, ACTION_ID)
# We expect 2 conflicts as stated in the fixture
zoo_keeper_client_with_conflict.set.call_count == 2
def test_cancel_happy_path(kubernetes_job_starter: DownloadDispatcher,
zoo_keeper_client_to_cancel,
mock_lock_manager: RWLockManager):
download_mgr = DownloadManager(zoo_keeper_client_to_cancel,
kubernetes_job_starter, mock_lock_manager,
data_set_to_path)
download_mgr.cancel(CLIENT_ID, ACTION_ID)
kubernetes_job_starter.cancel_all.assert_called_once_with(
CLIENT_ID, ACTION_ID)
zoo_keeper_client_to_cancel.set.assert_called_once_with(
SOME_PATH + "/" + DATASET_PATH, BIN_RESULT_CANCELED, ZK_VERSION)
def _test_fetch(zoo_keeper_client, enclosing_event, kubernetes_job_starter,
mock_lock_manager, some_data_set):
download_mgr = DownloadManager(
zoo_keeper_client,
kubernetes_job_starter,
mock_lock_manager,
data_set_to_path,
size_estimator=mock_size_estimator,
)
on_done = create_autospec(DownloadOnDone)
download_mgr.fetch(some_data_set, enclosing_event, on_done)
return on_done
def test_pass_through_estimator_error(
enclosing_event: BenchmarkEvent,
some_data_set: DownloadableContent,
failing_size_estimator: ContentSizeEstimator,
zoo_keeper_client: KazooClient,
kubernetes_job_starter: DownloadDispatcher,
mock_lock_manager: RWLockManager,
mock_lock: RWLock,
):
download_mgr = DownloadManager(zoo_keeper_client, kubernetes_job_starter,
mock_lock_manager, data_set_to_path,
failing_size_estimator)
on_done = create_autospec(DownloadOnDone)
download_mgr.fetch(some_data_set, enclosing_event, on_done)
_verify_failed_estimator(mock_lock, on_done, some_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 download_manager(zoo_keeper_client: KazooClient,
kubernetes_job_starter: DownloadDispatcher,
mock_lock_manager: RWLockManager) -> DownloadManager:
download_manager = DownloadManager(
zoo_keeper_client,
kubernetes_job_starter,
mock_lock_manager,
data_set_to_path,
size_estimator=mock_size_estimator,
)
return download_manager
def download_manager(zk_client: KazooClient,
k8s_dispatcher: KubernetesDispatcher):
locker = DistributedRWLockManager(zk_client, "it_locks")
download_mgr = DownloadManager(zk_client, k8s_dispatcher, locker)
download_mgr.start()
yield download_mgr
download_mgr.stop()