def test_get_environment():
with patch("prefect_saturn.core.Client", new=MockClient):
responses.add(**REGISTER_FLOW_RESPONSE())
responses.add(**BUILD_STORAGE_RESPONSE())
responses.add(**REGISTER_RUN_JOB_SPEC_RESPONSE(200))
integration = prefect_saturn.PrefectCloudIntegration(
prefect_cloud_project_name=TEST_PREFECT_PROJECT_NAME)
flow = TEST_FLOW.copy()
integration.register_flow_with_saturn(flow=flow)
environment = integration._get_environment(
cluster_kwargs={"n_workers": 3}, adapt_kwargs={})
assert isinstance(environment, KubernetesJobEnvironment)
assert environment.unique_job_name is True
env_args = environment._job_spec["spec"]["template"]["spec"][
"containers"][0]["args"]
assert len(env_args) == 1
assert env_args[0].startswith(
"source /home/jovyan/.saturn/start_wrapper.sh;")
env_cmd = environment._job_spec["spec"]["template"]["spec"][
"containers"][0]["command"]
assert env_cmd == ["/bin/bash", "-ec"]
assert environment.metadata["image"] == integration._saturn_image
assert len(environment.labels) == len(FLOW_LABELS)
for label in FLOW_LABELS:
assert label in environment.labels
def test_register_flow_with_saturn_does_everything():
with patch("prefect_saturn.core.Client", new=MockClient):
test_flow_id = str(random.randint(1, 500))
responses.add(**REGISTER_FLOW_RESPONSE(flow_id=test_flow_id))
responses.add(**BUILD_STORAGE_RESPONSE(flow_id=test_flow_id))
responses.add(
**REGISTER_RUN_JOB_SPEC_RESPONSE(200, flow_id=test_flow_id))
integration = prefect_saturn.PrefectCloudIntegration(
prefect_cloud_project_name=TEST_PREFECT_PROJECT_NAME)
# no details are stored just on initialization
assert integration._saturn_flow_id is None
assert integration._saturn_flow_version_id is None
assert integration._saturn_image is None
flow = TEST_FLOW.copy()
assert flow.storage is None
flow = integration.register_flow_with_saturn(flow=flow,
instance_size="large")
assert integration._saturn_flow_id == test_flow_id
assert integration._saturn_flow_version_id == TEST_FLOW_VERSION_ID
assert integration._saturn_image == TEST_IMAGE
assert isinstance(flow.storage, Webhook)
if RUN_CONFIG_AVAILABLE:
assert isinstance(flow.run_config, KubernetesRun)
assert isinstance(flow.executor, DaskExecutor)
else:
assert isinstance(flow.environment, KubernetesJobEnvironment)
assert isinstance(flow.environment.executor, DaskExecutor)
def test_executor_dask_kwargs():
with patch("prefect_saturn.core.Client", new=MockClient):
responses.add(**REGISTER_FLOW_RESPONSE())
responses.add(**BUILD_STORAGE_RESPONSE())
responses.add(**REGISTER_RUN_JOB_SPEC_RESPONSE(200))
integration = prefect_saturn.PrefectCloudIntegration(
prefect_cloud_project_name=TEST_PREFECT_PROJECT_NAME)
flow = TEST_FLOW.copy()
flow = integration.register_flow_with_saturn(
flow=flow,
dask_cluster_kwargs={
"n_workers": 8,
"autoclose": True
},
dask_adapt_kwargs={
"minimum": 3,
"maximum": 3
},
)
if RUN_CONFIG_AVAILABLE:
assert isinstance(flow.run_config, KubernetesRun)
executor = flow.executor
else:
assert isinstance(flow.environment, KubernetesJobEnvironment)
executor = flow.environment.executor
assert isinstance(executor, DaskExecutor)
assert executor.cluster_kwargs == {"n_workers": 8, "autoclose": True}
assert executor.adapt_kwargs == {"minimum": 3, "maximum": 3}
def test_initialize_raises_error_on_accessing_properties_if_not_registered():
integration = prefect_saturn.PrefectCloudIntegration(
prefect_cloud_project_name=TEST_PREFECT_PROJECT_NAME)
with raises(RuntimeError, match=prefect_saturn.Errors.NOT_REGISTERED):
integration.flow_id
with raises(RuntimeError, match=prefect_saturn.Errors.NOT_REGISTERED):
integration.flow_version_id
with raises(RuntimeError, match=prefect_saturn.Errors.NOT_REGISTERED):
integration.image
def test_hash_flow():
flow = TEST_FLOW.copy()
integration = prefect_saturn.PrefectCloudIntegration(
prefect_cloud_project_name=TEST_PREFECT_PROJECT_NAME)
with patch("prefect_saturn.core.Client", new=MockClient):
flow_hash = integration._hash_flow(flow)
assert isinstance(flow_hash, str) and len(flow_hash) > 0
# should be deterministic
flow_hash_again = integration._hash_flow(flow)
assert flow_hash == flow_hash_again
# should not be impacted by storage
flow.storage = Webhook(
build_request_kwargs={},
build_request_http_method="POST",
get_flow_request_kwargs={},
get_flow_request_http_method="GET",
)
assert flow_hash == integration._hash_flow(flow)
# should not be impacted by environment or run_config
if RUN_CONFIG_AVAILABLE:
flow.run_config = KubernetesRun()
elif KUBE_JOB_ENV_AVAILABLE:
flow.environment = KubernetesJobEnvironment()
assert flow_hash == integration._hash_flow(flow)
# should not change if you add a new task
@task
def goodbye_task():
logger = prefect.context.get("logger")
logger.info("adios")
flow.tasks = [hello_task, goodbye_task]
new_flow_hash = integration._hash_flow(flow)
assert isinstance(new_flow_hash, str) and len(new_flow_hash) > 0
assert new_flow_hash == flow_hash
# should change if flow name changes
flow.name = str(uuid.uuid4())
new_flow_hash = integration._hash_flow(flow)
assert new_flow_hash != flow_hash
# should change if project name changes
previous_flow_hash = new_flow_hash
integration.prefect_cloud_project_name = str(uuid.uuid4())
new_flow_hash = integration._hash_flow(flow)
assert isinstance(new_flow_hash, str) and len(new_flow_hash) > 0
assert new_flow_hash != previous_flow_hash
def test_add_environment_fails_if_id_not_recognized():
with patch("prefect_saturn.core.Client", new=MockClient):
responses.add(**REGISTER_FLOW_RESPONSE(flow_id=45))
responses.add(**REGISTER_RUN_JOB_SPEC_RESPONSE(404, flow_id=45))
integration = prefect_saturn.PrefectCloudIntegration(
prefect_cloud_project_name=TEST_PREFECT_PROJECT_NAME)
flow = TEST_FLOW.copy()
integration._set_flow_metadata(flow=flow, instance_size=None)
with raises(HTTPError, match="404 Client Error"):
integration._get_environment(cluster_kwargs={}, adapt_kwargs={})
def test_initialize():
project_name = "some-pr0ject"
integration = prefect_saturn.PrefectCloudIntegration(
prefect_cloud_project_name=project_name)
assert getattr(integration, "SATURN_TOKEN", None) is None
assert integration._session.headers[
"Authorization"] == f"token {os.environ['SATURN_TOKEN']}"
assert integration.prefect_cloud_project_name == project_name
assert integration._saturn_image is None
assert integration._saturn_flow_id is None
assert integration._saturn_flow_version_id is None
def test_store_flow_fails_if_validation_fails():
with patch("prefect_saturn.core.Client", new=MockClient):
responses.add(**REGISTER_FLOW_RESPONSE())
responses.add(**BUILD_STORAGE_FAILURE_RESPONSE(400))
responses.add(**REGISTER_RUN_JOB_SPEC_RESPONSE(200))
integration = prefect_saturn.PrefectCloudIntegration(
prefect_cloud_project_name=TEST_PREFECT_PROJECT_NAME)
flow = TEST_FLOW.copy()
flow = integration.register_flow_with_saturn(flow=flow)
with raises(HTTPError, match="Client Error"):
flow.storage.add_flow(flow)
flow.storage.build()
def test_register_flow_with_saturn_raises_error_on_failure():
with patch("prefect_saturn.core.Client", new=MockClient):
integration = prefect_saturn.PrefectCloudIntegration(
prefect_cloud_project_name=TEST_PREFECT_PROJECT_NAME)
responses.add(**REGISTER_FLOW_FAILURE_RESPONSE(500))
with raises(HTTPError, match="500 Server Error"):
integration.register_flow_with_saturn(flow=TEST_FLOW.copy())
failure_response = REGISTER_FLOW_FAILURE_RESPONSE(401)
failure_response["method_or_response"] = failure_response.pop("method")
responses.replace(**failure_response)
with raises(HTTPError, match="401 Client Error"):
integration.register_flow_with_saturn(flow=TEST_FLOW.copy())
def test_hash_flow_hash_changes_if_tenant_id_changes():
flow = TEST_FLOW.copy()
integration = prefect_saturn.PrefectCloudIntegration(
prefect_cloud_project_name=TEST_PREFECT_PROJECT_NAME)
with patch("prefect_saturn.core.Client", new=MockClient):
flow_hash = integration._hash_flow(flow)
assert isinstance(flow_hash, str) and len(flow_hash) > 0
class OtherMockClient:
def __init__(self):
self._active_tenant_id = "some-other-garbage"
with patch("prefect_saturn.core.Client", new=OtherMockClient):
new_flow_hash = integration._hash_flow(flow)
assert isinstance(new_flow_hash, str) and len(new_flow_hash) > 0
assert new_flow_hash != flow_hash
def test_get_storage():
with patch("prefect_saturn.core.Client", new=MockClient):
responses.add(**REGISTER_FLOW_RESPONSE())
responses.add(**BUILD_STORAGE_RESPONSE())
responses.add(**REGISTER_RUN_JOB_SPEC_RESPONSE(200))
integration = prefect_saturn.PrefectCloudIntegration(
prefect_cloud_project_name=TEST_PREFECT_PROJECT_NAME)
flow = TEST_FLOW.copy()
integration.register_flow_with_saturn(flow=flow)
storage = integration._get_storage()
assert isinstance(storage, Webhook)
assert storage.stored_as_script is False
assert storage.build_request_http_method == "POST"
assert storage.get_flow_request_http_method == "GET"
assert storage.build_request_kwargs["headers"][
"Content-Type"] == "application/octet-stream"
assert storage.get_flow_request_kwargs["headers"][
"Accept"] == "application/octet-stream"
def test_get_environment_fails_if_flow_not_registered():
integration = prefect_saturn.PrefectCloudIntegration(
prefect_cloud_project_name=TEST_PREFECT_PROJECT_NAME)
with raises(RuntimeError, match=prefect_saturn.Errors.NOT_REGISTERED):
integration._get_environment(cluster_kwargs={}, adapt_kwargs={})
def test_get_storage_fails_if_flow_not_registerd():
integration = prefect_saturn.PrefectCloudIntegration(
prefect_cloud_project_name=TEST_PREFECT_PROJECT_NAME)
with raises(RuntimeError, match=prefect_saturn.Errors.NOT_REGISTERED):
integration._get_storage()
def test_initialize_raises_error_on_missing_base_url(monkeypatch):
monkeypatch.delenv("BASE_URL")
with raises(RuntimeError,
match=prefect_saturn.Errors.missing_env_var("BASE_URL")):
prefect_saturn.PrefectCloudIntegration(
prefect_cloud_project_name=TEST_PREFECT_PROJECT_NAME)
def test_initialize_fails_on_extra_trailing_slash_in_base_url(monkeypatch):
monkeypatch.setenv("BASE_URL", "http://abc/")
with raises(RuntimeError, match=prefect_saturn.Errors.BASE_URL_NO_SLASH):
prefect_saturn.PrefectCloudIntegration("x")
