def test_docker_agent_deploy_flow_storage_raises(monkeypatch, api):
monkeypatch.setattr("prefect.agent.agent.Client", MagicMock())
agent = DockerAgent()
with pytest.raises(ValueError):
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage": Local().serialize(),
"id": "foo",
"name": "flow-name",
"environment": LocalEnvironment().serialize(),
"core_version": "0.13.0",
}),
"id":
"id",
"name":
"name",
"version":
"version",
}))
assert not api.pull.called
def test_query_flow_runs_ordered_by_start_time(monkeypatch, cloud_api):
dt1, dt2 = pendulum.now(), pendulum.now().add(hours=1)
gql_return = MagicMock(
return_value=MagicMock(
data=MagicMock(
get_runs_in_queue=MagicMock(flow_run_ids=["id"]),
flow_run=[
GraphQLResult({"id": "id2", "scheduled_start_time": str(dt2)}),
GraphQLResult({"id": "id", "scheduled_start_time": str(dt1)}),
],
)
)
)
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.agent.agent.Client", client)
agent = Agent()
flow_runs = agent.query_flow_runs()
assert flow_runs == [
GraphQLResult({"id": "id", "scheduled_start_time": str(dt1)}),
GraphQLResult({"id": "id2", "scheduled_start_time": str(dt2)}),
]
def test_docker_agent_deploy_flow_no_registry_does_not_pull(api):
agent = DockerAgent()
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"name":
"flow-name",
"storage":
Docker(registry_url="", image_name="name",
image_tag="tag").serialize(),
"core_version":
"0.13.0",
}),
"id":
"id",
"name":
"name",
}))
assert not api.pull.called
assert api.create_container.called
assert api.start.called
def test_docker_agent_networks(api):
api.create_networking_config.return_value = {
"test-network-1": "config1",
"test-network-2": "config2",
}
agent = DockerAgent(networks=["test-network-1", "test-network-2"])
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"name":
"flow-name",
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
LocalEnvironment().serialize(),
"core_version":
"0.13.0",
}),
"id":
"id",
"name":
"name",
}))
assert "test-network-1" in agent.networks
assert "test-network-2" in agent.networks
assert agent.network is None
args, kwargs = api.create_container.call_args
assert kwargs["networking_config"] == {
"test-network-1": "config1",
"test-network-2": "config2",
}
def test_k8s_agent_deploy_flow_raises(monkeypatch, cloud_api):
k8s_config = MagicMock()
monkeypatch.setattr("kubernetes.config", k8s_config)
batch_client = MagicMock()
batch_client.create_namespaced_job.return_value = {}
monkeypatch.setattr("kubernetes.client.BatchV1Api",
MagicMock(retrurn_value=batch_client))
agent = KubernetesAgent()
with pytest.raises(ValueError):
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage": Local().serialize(),
"id": "id",
"environment": LocalEnvironment().serialize(),
"core_version": "0.13.0",
}),
"id":
"id",
}))
assert not agent.batch_client.create_namespaced_job.called
def test_k8s_agent_replace_yaml_responds_to_logging_config(
monkeypatch, runner_token, flag):
k8s_config = MagicMock()
monkeypatch.setattr("kubernetes.config", k8s_config)
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"id":
"new_id",
}),
"id":
"id",
"name":
"name",
})
agent = KubernetesAgent(no_cloud_logs=flag)
job = agent.replace_job_spec_yaml(flow_run)
env = job["spec"]["template"]["spec"]["containers"][0]["env"]
assert env[6]["value"] == str(not flag).lower()
def test_docker_agent_deploy_flow_show_flow_logs(monkeypatch, runner_token):
process = MagicMock()
monkeypatch.setattr("multiprocessing.Process", process)
api = MagicMock()
api.ping.return_value = True
api.create_container.return_value = {"Id": "container_id"}
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
agent = DockerAgent(show_flow_logs=True)
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
LocalEnvironment().serialize(),
}),
"id":
"id",
"name":
"name",
}))
process.assert_called_with(target=agent.stream_container_logs,
kwargs={"container_id": "container_id"})
assert len(agent.processes) == 1
assert api.create_container.called
assert api.start.called
def test_update_states_passes_no_task_runs(monkeypatch, cloud_api):
gql_return = MagicMock(return_value=MagicMock(
data=MagicMock(set_flow_run_state=None, set_task_run_state=None)))
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.agent.agent.Client", client)
agent = Agent()
assert not agent.update_state(
flow_run=GraphQLResult({
"id": "id",
"serialized_state": Scheduled().serialize(),
"version": 1,
"task_runs": [],
}))
def test_mark_flow_as_failed(monkeypatch, cloud_api):
agent = Agent()
agent.client = MagicMock()
agent._mark_flow_as_failed(
flow_run=GraphQLResult({
"id": "id",
"serialized_state": Scheduled().serialize(),
"version": 1,
"task_runs": [],
}),
message="foo",
)
agent.client.set_flow_run_state.assert_called_with(
flow_run_id="id", version=1, state=Failed(message="foo"))
def get_run_task_kwargs(
self, run_config, tenant_id: str = None, taskdef=None, **kwargs
):
agent = ECSAgent(**kwargs)
if tenant_id:
agent.client._get_auth_tenant = MagicMock(return_value=tenant_id)
flow_run = GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Local().serialize(),
"id": "flow-id",
"version": 1,
"name": "Test Flow",
"core_version": "0.13.0",
}
),
"run_config": run_config.serialize(),
"id": "flow-run-id",
}
)
if taskdef is None:
taskdef = agent.generate_task_definition(flow_run, run_config)
return agent.get_run_task_kwargs(flow_run, run_config, taskdef)
def test_delete_key_value_raises_if_key_not_found(self, monkeypatch,
cloud_api):
Client = MagicMock()
Client().graphql.return_value = GraphQLResult(
data=dict(key_value=[], ))
monkeypatch.setattr("prefect.backend.kv_store.Client", Client)
with pytest.raises(ValueError):
delete_key(key="foo")
Client.return_value.graphql.assert_called_with(query={
"query": {
'key_value(where: { key: { _eq: "foo" } })': {"id"}
}
})
def test_local_agent_deploy_keep_existing_python_path(monkeypatch,
runner_token):
monkeypatch.setenv("PYTHONPATH", "cool:python:path")
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent(import_paths=["paths"])
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage": Local(directory="test").serialize(),
"id": "foo"
}),
"id":
"id",
}))
python_path = popen.call_args[1]["env"]["PYTHONPATH"]
assert popen.called
assert len(agent.processes) == 1
assert "cool:python:path" in python_path
assert "paths" in python_path
def test_local_agent_deploy_storage_fails_none(monkeypatch, runner_token):
client = MagicMock()
set_state = MagicMock()
client.return_value.set_flow_run_state = set_state
monkeypatch.setattr("prefect.agent.agent.Client", client)
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent()
with pytest.raises(ValidationError):
agent.deploy_flow(
flow_run=GraphQLResult({
"flow": GraphQLResult({"storage": None}),
"id": "id",
"version": 1
}))
assert not popen.called
assert len(agent.processes) == 0
assert client.called
def test_deploy_flow_run_sleeps_until_start_time(monkeypatch, cloud_api):
gql_return = MagicMock(
return_value=MagicMock(data=MagicMock(write_run_logs=MagicMock(success=True)))
)
client = MagicMock()
client.return_value.write_run_logs = gql_return
monkeypatch.setattr("prefect.agent.agent.Client", MagicMock(return_value=client))
sleep = MagicMock()
monkeypatch.setattr("time.sleep", sleep)
dt = pendulum.now()
agent = Agent()
agent.deploy_flow = MagicMock()
agent._deploy_flow_run(
flow_run=GraphQLResult(
{
"id": "id",
"serialized_state": Scheduled().serialize(),
"scheduled_start_time": str(dt.add(seconds=10)),
"version": 1,
"task_runs": [
GraphQLResult(
{
"id": "id",
"version": 1,
"serialized_state": Scheduled().serialize(),
}
)
],
}
)
)
sleep_time = sleep.call_args[0][0]
assert 10 >= sleep_time > 9
agent.deploy_flow.assert_called_once()
def test_k8s_agent_removes_yaml_no_volume(monkeypatch, cloud_api):
flow_run = GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Docker(
registry_url="test", image_name="name", image_tag="tag"
).serialize(),
"environment": LocalEnvironment().serialize(),
"id": "id",
"core_version": "0.13.0",
}
),
"id": "id",
}
)
agent = KubernetesAgent()
job = agent.generate_job_spec_from_environment(flow_run, image="test/name:tag")
assert not job["spec"]["template"]["spec"].get("volumes", None)
assert not job["spec"]["template"]["spec"]["containers"][0].get(
"volumeMounts", None
)
def test_local_agent_deploy_processes_webhook_storage(monkeypatch,
runner_token):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent()
webhook = Webhook(
build_request_kwargs={"url": "test-service/upload"},
build_request_http_method="POST",
get_flow_request_kwargs={"url": "test-service/download"},
get_flow_request_http_method="GET",
)
agent.deploy_flow(flow_run=GraphQLResult(
{
"flow": GraphQLResult({
"storage": webhook.serialize(),
"id": "foo"
}),
"id": "id",
}))
assert popen.called
assert len(agent.processes) == 1
def test_populate_env_vars(cloud_api):
with set_temporary_config({
"cloud.api": "api",
"logging.log_to_cloud": True,
"cloud.agent.auth_token": "token",
}):
agent = LocalAgent()
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}))
expected_vars = {
"PREFECT__CLOUD__API":
"api",
"PREFECT__CLOUD__AUTH_TOKEN":
"token",
"PREFECT__CLOUD__AGENT__LABELS":
str([
socket.gethostname(),
"azure-flow-storage",
"gcs-flow-storage",
"s3-flow-storage",
"github-flow-storage",
"webhook-flow-storage",
]),
"PREFECT__CONTEXT__FLOW_RUN_ID":
"id",
"PREFECT__CONTEXT__FLOW_ID":
"foo",
"PREFECT__CLOUD__USE_LOCAL_SECRETS":
"false",
"PREFECT__LOGGING__LOG_TO_CLOUD":
"true",
"PREFECT__LOGGING__LEVEL":
"DEBUG",
"PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS":
"prefect.engine.cloud.CloudFlowRunner",
"PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS":
"prefect.engine.cloud.CloudTaskRunner",
}
assert env_vars == expected_vars
def test_populate_env_vars_is_responsive_to_logging_config(
monkeypatch, cloud_api, flag
):
api = MagicMock()
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
with set_temporary_config({"cloud.agent.auth_token": "token", "cloud.api": "api"}):
agent = DockerAgent(labels=["42", "marvin"], no_cloud_logs=flag)
env_vars = agent.populate_env_vars(
GraphQLResult({"id": "id", "name": "name", "flow": {"id": "foo"}})
)
assert env_vars["PREFECT__LOGGING__LOG_TO_CLOUD"] == str(not flag).lower()
def test_populate_env_vars_sets_log_to_cloud(flag, api, config_with_token):
agent = DockerAgent(labels=["42", "marvin"], no_cloud_logs=flag)
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}), "test-image")
assert env_vars["PREFECT__CLOUD__SEND_FLOW_RUN_LOGS"] == str(
not flag).lower()
# Backwards compatibility variable for containers on Prefect <0.15.0
assert env_vars["PREFECT__LOGGING__LOG_TO_CLOUD"] == str(not flag).lower()
def test_populate_env_vars_keep_existing_python_path(monkeypatch):
monkeypatch.setenv("PYTHONPATH", "cool:python:path")
agent = LocalAgent(import_paths=["paths"])
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}))
python_path = env_vars["PYTHONPATH"]
assert "cool:python:path" in python_path
assert "paths" in python_path
def test_local_agent_responds_to_logging_config(runner_token, flag):
with set_temporary_config({
"cloud.agent.auth_token": "TEST_TOKEN",
"logging.log_to_cloud": flag
}):
agent = LocalAgent()
assert agent.log_to_cloud is flag
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}))
assert env_vars["PREFECT__LOGGING__LOG_TO_CLOUD"] == str(flag).lower()
def test_populate_env_vars_includes_agent_labels(monkeypatch, cloud_api):
api = MagicMock()
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
with set_temporary_config({
"cloud.agent.auth_token": "token",
"cloud.api": "api",
"logging.log_to_cloud": True,
}):
agent = DockerAgent(labels=["42", "marvin"])
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}))
expected_vars = {
"PREFECT__CLOUD__API":
"api",
"PREFECT__CLOUD__AGENT__LABELS":
"['42', 'marvin']",
"PREFECT__CLOUD__AUTH_TOKEN":
"token",
"PREFECT__CONTEXT__FLOW_RUN_ID":
"id",
"PREFECT__CONTEXT__FLOW_ID":
"foo",
"PREFECT__CLOUD__USE_LOCAL_SECRETS":
"false",
"PREFECT__LOGGING__LOG_TO_CLOUD":
"true",
"PREFECT__LOGGING__LEVEL":
"INFO",
"PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS":
"prefect.engine.cloud.CloudFlowRunner",
"PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS":
"prefect.engine.cloud.CloudTaskRunner",
}
assert env_vars == expected_vars
def test_populate_env_vars(api, backend):
agent = DockerAgent()
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}), "test-image")
if backend == "server":
cloud_api = "http://host.docker.internal:4200"
else:
cloud_api = prefect.config.cloud.api
expected_vars = {
"PREFECT__BACKEND":
backend,
"PREFECT__CLOUD__API":
cloud_api,
"PREFECT__CLOUD__AUTH_TOKEN":
"TEST_TOKEN",
"PREFECT__CLOUD__AGENT__LABELS":
"[]",
"PREFECT__CONTEXT__FLOW_RUN_ID":
"id",
"PREFECT__CONTEXT__FLOW_ID":
"foo",
"PREFECT__CONTEXT__IMAGE":
"test-image",
"PREFECT__CLOUD__USE_LOCAL_SECRETS":
"false",
"PREFECT__CLOUD__SEND_FLOW_RUN_LOGS":
"true",
"PREFECT__LOGGING__LOG_TO_CLOUD":
"true",
"PREFECT__LOGGING__LEVEL":
"INFO",
"PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS":
"prefect.engine.cloud.CloudFlowRunner",
"PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS":
"prefect.engine.cloud.CloudTaskRunner",
}
assert env_vars == expected_vars
def test_mark_failed(monkeypatch, runner_token, cloud_api):
gql_return = MagicMock(return_value=MagicMock(
data=MagicMock(set_flow_run_state=None, set_task_run_state=None)))
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.agent.agent.Client", client)
agent = Agent()
assert not agent.mark_failed(
flow_run=GraphQLResult({
"id": "id",
"serialized_state": Scheduled().serialize(),
"version": 1,
"task_runs": [],
}),
exc=Exception(),
)
def test_get_ready_flow_runs(monkeypatch, cloud_api):
dt = pendulum.now()
gql_return = MagicMock(
return_value=MagicMock(
data=MagicMock(
get_runs_in_queue=MagicMock(flow_run_ids=["id"]),
flow_run=[GraphQLResult({"id": "id", "scheduled_start_time": str(dt)})],
)
)
)
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.agent.agent.Client", client)
agent = Agent()
flow_runs = agent._get_ready_flow_runs()
assert flow_runs == {"id"}
def test_populate_env_vars(monkeypatch, backend):
agent = LocalAgent()
# The python path may be a single item and we want to ensure the correct separator
# is added so we will ensure PYTHONPATH has an item in it to start
if not os.environ.get("PYTHONPATH", ""):
monkeypatch.setenv("PYTHONPATH", "foobar")
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}))
expected = os.environ.copy()
expected.update({
"PYTHONPATH":
os.getcwd() + os.pathsep + expected.get("PYTHONPATH", ""),
"PREFECT__BACKEND":
backend,
"PREFECT__CLOUD__API":
prefect.config.cloud.api,
"PREFECT__CLOUD__AUTH_TOKEN":
"TEST_TOKEN",
"PREFECT__CLOUD__AGENT__LABELS":
str(DEFAULT_AGENT_LABELS),
"PREFECT__CONTEXT__FLOW_RUN_ID":
"id",
"PREFECT__CONTEXT__FLOW_ID":
"foo",
"PREFECT__CLOUD__USE_LOCAL_SECRETS":
"false",
"PREFECT__LOGGING__LOG_TO_CLOUD":
"true",
"PREFECT__LOGGING__LEVEL":
"INFO",
"PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS":
"prefect.engine.cloud.CloudFlowRunner",
"PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS":
"prefect.engine.cloud.CloudTaskRunner",
})
assert env_vars == expected
def test_update_states_passes_no_task_runs(monkeypatch):
with set_temporary_config({"cloud.agent.auth_token": "token"}):
gql_return = MagicMock(return_value=MagicMock(
data=MagicMock(set_flow_run_state=None, set_task_run_state=None)))
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.agent.agent.Client", client)
agent = Agent()
assert not agent.update_states(flow_runs=[
GraphQLResult({
"id": "id",
"serialized_state": Scheduled().serialize(),
"version": 1,
"task_runs": [],
})
])
def graphql(
self,
query: Any,
raise_on_error: bool = True,
headers: Dict[str, str] = None,
variables: Dict[str, JSONLike] = None,
token: str = None,
) -> GraphQLResult:
"""
Convenience function for running queries against the Prefect GraphQL API
Args:
- query (Any): A representation of a graphql query to be executed. It will be
parsed by prefect.utilities.graphql.parse_graphql().
- raise_on_error (bool): if True, a `ClientError` will be raised if the GraphQL
returns any `errors`.
- headers (dict): any additional headers that should be passed as part of the
request
- variables (dict): Variables to be filled into a query with the key being
equivalent to the variables that are accepted by the query
- token (str): an auth token. If not supplied, the `client.access_token` is used.
Returns:
- dict: Data returned from the GraphQL query
Raises:
- ClientError if there are errors raised by the GraphQL mutation
"""
result = self.post(
path="",
server=self.api_server,
headers=headers,
params=dict(query=parse_graphql(query),
variables=json.dumps(variables)),
token=token,
)
if raise_on_error and "errors" in result:
if "UNAUTHENTICATED" in str(result["errors"]):
raise AuthorizationError(result["errors"])
elif "Malformed Authorization header" in str(result["errors"]):
raise AuthorizationError(result["errors"])
raise ClientError(result["errors"])
else:
return GraphQLResult(result) # type: ignore
def test_prefect_logging_level_override_logic(config, agent_env_vars,
run_config_env_vars,
expected_logging_level, tmpdir):
with set_temporary_config(config):
agent = LocalAgent(env_vars=agent_env_vars)
run = LocalRun(working_dir=str(tmpdir), env=run_config_env_vars)
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
},
"run_config": run.serialize(),
}),
run,
)
assert env_vars["PREFECT__LOGGING__LEVEL"] == expected_logging_level
def test_populate_env_vars_uses_user_provided_env_vars(cloud_api):
with set_temporary_config({
"cloud.api": "api",
"logging.log_to_cloud": True,
"cloud.agent.auth_token": "token",
}):
agent = LocalAgent(env_vars=dict(AUTH_THING="foo"))
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}))
assert env_vars["AUTH_THING"] == "foo"
