def test_populate_worker_pod_yaml(): environment = CloudEnvironment() file_path = os.path.dirname( prefect.environments.execution.cloud.environment.__file__) with open(path.join(file_path, "worker_pod.yaml")) as pod_file: pod = yaml.safe_load(pod_file) with set_temporary_config({ "cloud.graphql": "gql_test", "cloud.auth_token": "auth_test" }): with prefect.context(flow_run_id="id_test", image="my_image"): yaml_obj = environment._populate_worker_pod_yaml(yaml_obj=pod) assert yaml_obj["metadata"]["labels"][ "identifier"] == environment.identifier_label assert yaml_obj["metadata"]["labels"]["flow_run_id"] == "id_test" env = yaml_obj["spec"]["containers"][0]["env"] assert env[0]["value"] == "gql_test" assert env[1]["value"] == "auth_test" assert env[2]["value"] == "id_test" assert yaml_obj["spec"]["containers"][0]["image"] == "my_image"
def test_run_flow(monkeypatch): environment = CloudEnvironment() flow_runner = MagicMock() monkeypatch.setattr( "prefect.engine.get_default_flow_runner_class", MagicMock(return_value=flow_runner), ) monkeypatch.setattr( "prefect.environments.execution.cloud.environment.sys.exit", MagicMock()) kube_cluster = MagicMock() monkeypatch.setattr("dask_kubernetes.KubeCluster", kube_cluster) with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "flow_env.prefect"), "w+") as env: flow = prefect.Flow("test") flow_path = os.path.join(directory, "flow_env.prefect") with open(flow_path, "wb") as f: cloudpickle.dump(flow, f) with set_temporary_config({"cloud.auth_token": "test"}): with prefect.context(flow_file_path=os.path.join( directory, "flow_env.prefect")): environment.run_flow() assert flow_runner.call_args[1]["flow"].name == "test"
def test_create_flow_run_job_fails_outside_cluster(): environment = CloudEnvironment() with pytest.raises(EnvironmentError): with set_temporary_config({"cloud.auth_token": "test"}): environment.create_flow_run_job(docker_name="test1/test2:test3", flow_file_path="test4")
def test_execute(monkeypatch): environment = CloudEnvironment() storage = Docker(registry_url="test1", image_name="test2", image_tag="test3") create_flow_run = MagicMock() monkeypatch.setattr( "prefect.environments.CloudEnvironment.create_flow_run_job", create_flow_run ) environment.execute(storage=storage, flow_location="") assert create_flow_run.call_args[1]["docker_name"] == "test1/test2:test3"
def test_populate_job_yaml(): environment = CloudEnvironment() file_path = os.path.dirname( prefect.environments.execution.cloud.environment.__file__ ) with open(path.join(file_path, "job.yaml")) as job_file: job = yaml.safe_load(job_file) with set_temporary_config( { "cloud.graphql": "gql_test", "cloud.log": "log_test", "cloud.result_handler": "rh_test", "cloud.auth_token": "auth_test", } ): with prefect.context(flow_run_id="id_test", namespace="namespace_test"): yaml_obj = environment._populate_job_yaml( yaml_obj=job, docker_name="test1/test2:test3", flow_file_path="test4" ) assert yaml_obj["metadata"]["name"] == "prefect-dask-job-{}".format( environment.identifier_label ) assert yaml_obj["metadata"]["labels"]["identifier"] == environment.identifier_label assert yaml_obj["metadata"]["labels"]["flow_run_id"] == "id_test" assert ( yaml_obj["spec"]["template"]["metadata"]["labels"]["identifier"] == environment.identifier_label ) env = yaml_obj["spec"]["template"]["spec"]["containers"][0]["env"] assert env[0]["value"] == "gql_test" assert env[1]["value"] == "log_test" assert env[2]["value"] == "rh_test" assert env[3]["value"] == "auth_test" assert env[4]["value"] == "id_test" assert env[5]["value"] == "namespace_test" assert env[6]["value"] == "test1/test2:test3" assert env[7]["value"] == "test4" assert ( yaml_obj["spec"]["template"]["spec"]["containers"][0]["image"] == "test1/test2:test3" )
def test_create_flow_run_job(monkeypatch): environment = CloudEnvironment() config = MagicMock() monkeypatch.setattr("kubernetes.config", config) batchv1 = MagicMock() monkeypatch.setattr("kubernetes.client", MagicMock(BatchV1Api=MagicMock(return_value=batchv1))) with set_temporary_config({"cloud.auth_token": "test"}): environment.create_flow_run_job(docker_name="test1/test2:test3", flow_file_path="test4") assert (batchv1.create_namespaced_job.call_args[1]["body"]["apiVersion"] == "batch/v1")
def test_populate_worker_pod_yaml_with_private_registry(): environment = CloudEnvironment(private_registry=True) file_path = os.path.dirname( prefect.environments.execution.cloud.environment.__file__) with open(path.join(file_path, "worker_pod.yaml")) as pod_file: pod = yaml.safe_load(pod_file) with set_temporary_config({ "cloud.graphql": "gql_test", "cloud.auth_token": "auth_test" }): with prefect.context(flow_run_id="id_test", image="my_image", namespace="foo-man"): yaml_obj = environment._populate_worker_pod_yaml(yaml_obj=pod) yaml_obj["spec"]["imagePullSecrets"][0] == dict(name="foo-man-docker")
def test_setup_doesnt_pass_if_private_registry(monkeypatch): environment = CloudEnvironment(private_registry=True) config = MagicMock() monkeypatch.setattr("kubernetes.config", config) v1 = MagicMock() v1.list_namespaced_secret.return_value = MagicMock(items=[]) monkeypatch.setattr("kubernetes.client", MagicMock(CoreV1Api=MagicMock(return_value=v1))) create_secret = MagicMock() monkeypatch.setattr( "prefect.environments.CloudEnvironment._create_namespaced_secret", create_secret) with set_temporary_config({"cloud.auth_token": "test"}): environment.setup(storage=Docker()) assert create_secret.called
def test_create_secret_isnt_called_if_exists(monkeypatch): environment = CloudEnvironment(private_registry=True) config = MagicMock() monkeypatch.setattr("kubernetes.config", config) secret = MagicMock() secret.metadata.name = "foo-docker" v1 = MagicMock() v1.list_namespaced_secret.return_value = MagicMock(items=[secret]) monkeypatch.setattr("kubernetes.client", MagicMock(CoreV1Api=MagicMock(return_value=v1))) create_secret = MagicMock() monkeypatch.setattr( "prefect.environments.CloudEnvironment._create_namespaced_secret", create_secret) with set_temporary_config({"cloud.auth_token": "test"}): with prefect.context(namespace="foo"): environment.setup(storage=Docker()) assert not create_secret.called
def test_run_flow(monkeypatch): environment = CloudEnvironment() flow_runner = MagicMock() monkeypatch.setattr("prefect.engine.FlowRunner", flow_runner) kube_cluster = MagicMock() monkeypatch.setattr("dask_kubernetes.KubeCluster", kube_cluster) with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "flow_env.prefect"), "w+") as env: flow = prefect.Flow("test") flow_path = os.path.join(directory, "flow_env.prefect") with open(flow_path, "w") as f: json.dump(flow.serialize(), f) with set_temporary_config({"cloud.auth_token": "test"}): with prefect.context(flow_file_path=os.path.join( directory, "flow_env.prefect")): environment.run_flow() assert flow_runner.call_args[1]["flow"].name == "test"
def test_execute_storage_missing_fields(): environment = CloudEnvironment() with pytest.raises(ValueError): environment.execute(storage=Docker(), flow_location="")
def test_execute_improper_storage(): environment = CloudEnvironment() with pytest.raises(TypeError): environment.execute(storage=Memory(), flow_location="")
def test_setup_cloud_environment_passes(): environment = CloudEnvironment() environment.setup(storage=Docker()) assert environment
def test_create_cloud_environment_identifier_label(): environment = CloudEnvironment() assert environment.identifier_label
def test_create_cloud_environment(): environment = CloudEnvironment() assert environment assert environment.private_registry is False assert environment.docker_secret is None
def test_create_cloud_environment(): environment = CloudEnvironment() assert environment