def test_local_agent_deploy_flows(monkeypatch):
api = MagicMock()
api.ping.return_value = True
api.create_container.return_value = {"Id": "container_id"}
monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient",
MagicMock(return_value=api))
with set_temporary_config({"cloud.agent.auth_token": "token"}):
agent = LocalAgent()
agent.deploy_flows(flow_runs=[
GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(
registry_url="test",
image_name="name",
image_tag="tag",
).serialize()
}),
"id":
"id",
})
])
assert api.pull.called
assert api.create_container.called
assert api.start.called
assert (api.create_container.call_args[1]["command"] ==
"prefect execute cloud-flow")
assert api.start.call_args[1]["container"] == "container_id"
def test_local_agent_uses_ip_if_dockerdesktop_hostname(monkeypatch):
monkeypatch.setattr("socket.gethostname",
MagicMock(return_value="docker-desktop"))
monkeypatch.setattr("socket.gethostbyname", MagicMock(return_value="IP"))
agent = LocalAgent()
assert "IP" in agent.labels
def test_local_agent_deploy_no_existing_python_path(monkeypatch, cloud_api):
monkeypatch.delenv("PYTHONPATH", raising=False)
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",
"core_version": "0.13.0",
}
),
"id": "id",
}
)
)
assert popen.called
assert len(agent.processes) == 1
assert "paths" in popen.call_args[1]["env"]["PYTHONPATH"]
def test_local_agent_init(monkeypatch, runner_token):
api = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", api)
agent = LocalAgent()
assert agent
assert agent.name == "agent"
def test_populate_env_vars(monkeypatch):
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__CLOUD__API": "https://api.prefect.io",
"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_local_agent_deploy_storage_raises_not_supported_storage(
monkeypatch, cloud_api
):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent()
with pytest.raises(ValueError):
agent.deploy_flow(
flow_run=GraphQLResult(
{
"id": "id",
"flow": {
"storage": Docker().serialize(),
"id": "foo",
"core_version": "0.13.0",
},
},
)
)
assert not popen.called
assert len(agent.processes) == 0
def test_local_agent_deploy_run_config_missing_working_dir(monkeypatch, tmpdir):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
working_dir = str(tmpdir.join("missing"))
agent = LocalAgent()
with pytest.raises(ValueError, match="nonexistent `working_dir`"):
agent.deploy_flow(
flow_run=GraphQLResult(
{
"id": "id",
"flow": {
"storage": Local().serialize(),
"id": "foo",
"core_version": "0.13.0",
},
"run_config": LocalRun(working_dir=working_dir).serialize(),
},
)
)
assert not popen.called
assert not agent.processes
def test_local_agent_deploy_run_config_working_dir(monkeypatch, working_dir, tmpdir):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
if working_dir is not None:
working_dir = str(tmpdir)
agent = LocalAgent()
agent.deploy_flow(
flow_run=GraphQLResult(
{
"id": "id",
"flow": {
"storage": Local().serialize(),
"id": "foo",
"core_version": "0.13.0",
},
"run_config": LocalRun(working_dir=working_dir).serialize(),
},
)
)
assert popen.called
assert len(agent.processes) == 1
assert popen.call_args[1]["cwd"] == working_dir
def test_local_agent_deploy_unsupported_run_config(monkeypatch):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent()
with pytest.raises(
TypeError,
match="`run_config` of type `KubernetesRun`, only `LocalRun` is supported",
):
agent.deploy_flow(
flow_run=GraphQLResult(
{
"id": "id",
"flow": {
"storage": Local().serialize(),
"id": "foo",
"core_version": "0.13.0",
},
"run_config": KubernetesRun().serialize(),
},
)
)
assert not popen.called
assert len(agent.processes) == 0
def test_local_agent_deploy_flows_no_registry_does_not_pull(
monkeypatch, runner_token):
api = MagicMock()
api.ping.return_value = True
api.create_container.return_value = {"Id": "container_id"}
monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient",
MagicMock(return_value=api))
agent = LocalAgent()
agent.deploy_flows(flow_runs=[
GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="", image_name="name",
image_tag="tag").serialize()
}),
"id":
"id",
})
])
assert not api.pull.called
assert api.create_container.called
assert api.start.called
def test_local_agent_deploy_storage_fails_none(monkeypatch):
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": "foo",
"core_version": "0.13.0"
}),
"id":
"id",
"version":
1,
}))
assert not popen.called
assert len(agent.processes) == 0
assert client.called
def test_populate_env_vars(runner_token):
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"}))
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",
]
),
"PREFECT__CONTEXT__FLOW_RUN_ID": "id",
"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(monkeypatch, runner_token):
api = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", api)
with set_temporary_config({
"cloud.agent.auth_token": "token",
"cloud.api": "api"
}):
agent = LocalAgent()
env_vars = agent.populate_env_vars(GraphQLResult({"id": "id"}))
expected_vars = {
"PREFECT__CLOUD__API":
"api",
"PREFECT__CLOUD__AUTH_TOKEN":
"token",
"PREFECT__CONTEXT__FLOW_RUN_ID":
"id",
"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_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",
"core_version": "0.13.0",
}),
"id":
"id",
}))
assert popen.called
assert len(agent.processes) == 1
def test_local_agent_config_options(cloud_api):
with set_temporary_config({
"cloud.agent.auth_token": "TEST_TOKEN",
"logging.log_to_cloud": True
}):
agent = LocalAgent(
name="test",
labels=["test_label"],
import_paths=["test_path"],
hostname_label=False,
)
assert agent.name == "test"
assert agent.client.get_auth_token() == "TEST_TOKEN"
assert agent.logger
assert agent.log_to_cloud is True
assert agent.processes == set()
assert agent.import_paths == ["test_path"]
assert set(agent.labels) == {
"azure-flow-storage",
"s3-flow-storage",
"gcs-flow-storage",
"github-flow-storage",
"webhook-flow-storage",
"test_label",
}
def test_populate_env_vars_sets_log_to_cloud(flag):
agent = LocalAgent(no_cloud_logs=flag)
assert agent.log_to_cloud is not 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_from_run_config(tmpdir):
agent = LocalAgent(env_vars={"KEY1": "VAL1", "KEY2": "VAL2"})
working_dir = str(tmpdir)
run = LocalRun(
env={
"KEY2": "OVERRIDE",
"PREFECT__LOGGING__LEVEL": "TEST"
},
working_dir=working_dir,
)
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
},
"run_config": run.serialize(),
}),
run,
)
assert env_vars["KEY1"] == "VAL1"
assert env_vars["KEY2"] == "OVERRIDE"
assert env_vars["PREFECT__LOGGING__LEVEL"] == "TEST"
assert working_dir in env_vars["PYTHONPATH"]
def test_populate_env_vars_includes_agent_labels():
agent = LocalAgent(labels=["42", "marvin"])
env_vars = agent.populate_env_vars(
GraphQLResult({"id": "id", "name": "name", "flow": {"id": "foo"}})
)
expected = str(["42", "marvin"] + DEFAULT_AGENT_LABELS)
assert env_vars["PREFECT__CLOUD__AGENT__LABELS"] == expected
def test_populate_env_vars_removes_none_values():
agent = LocalAgent(env_vars=dict(MISSING_VAR=None))
env_vars = agent.populate_env_vars(
GraphQLResult({"id": "id", "name": "name", "flow": {"id": "foo"}})
)
assert "MISSING_VAR" not in env_vars
def test_local_agent_responds_to_logging_config(runner_token, flag):
with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = LocalAgent(no_cloud_logs=flag)
assert agent.log_to_cloud is not 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_no_existing_python_path(monkeypatch):
monkeypatch.delenv("PYTHONPATH", raising=False)
agent = LocalAgent(import_paths=["paths"])
env_vars = agent.populate_env_vars(
GraphQLResult({"id": "id", "name": "name", "flow": {"id": "foo"}})
)
assert "paths" in env_vars["PYTHONPATH"]
def test_populate_env_vars_from_agent_config():
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"
def test_environment_has_agent_token_from_config():
"""Check that the API token is passed through from the config via environ"""
with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = LocalAgent()
env = agent.populate_env_vars(TEST_FLOW_RUN_DATA)
assert env["PREFECT__CLOUD__AUTH_TOKEN"] == "TEST_TOKEN"
def test_local_agent_ping(monkeypatch, runner_token):
api = MagicMock()
api.ping.return_value = True
monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient",
MagicMock(return_value=api))
agent = LocalAgent()
assert api.ping.called
def test_local_agent_heartbeat(
monkeypatch, cloud_api, returncode, show_flow_logs, logs
):
popen = MockPopen()
# expect a process to be called with the following command (with specified behavior)
popen.set_command(
"prefect execute flow-run",
stdout=b"awesome output!",
stderr=b"blerg, eRroR!",
returncode=returncode,
poll_count=2,
)
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent(import_paths=["paths"], show_flow_logs=show_flow_logs)
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Local(directory="test").serialize(),
"id": "foo",
"core_version": "0.13.0",
}
),
"id": "id",
}
)
)
process = list(agent.processes)[0]
process_call = process.root_call
with LogCapture() as logcap:
agent.heartbeat()
agent.heartbeat()
agent.heartbeat()
# ensure the expected logs exist (or the absense of logs)
if logs:
logcap.check(*logs)
else:
logcap.check()
# ensure the process was opened and was polled
compare(
popen.all_calls,
expected=[
process_call,
process_call.poll(),
process_call.poll(),
process_call.poll(),
],
)
# the heartbeat should stop tracking upon exit
compare(process.returncode, returncode)
assert len(agent.processes) == 0
def test_local_agent_ping(monkeypatch):
api = MagicMock()
api.ping.return_value = True
monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient",
MagicMock(return_value=api))
with set_temporary_config({"cloud.agent.auth_token": "token"}):
agent = LocalAgent()
assert api.ping.called
def test_local_agent_config_no_storage_labels():
agent = LocalAgent(
labels=["test_label"],
storage_labels=False,
)
assert set(agent.labels) == {
socket.gethostname(),
"test_label",
}
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(TEST_FLOW_RUN_DATA)
python_path = env_vars["PYTHONPATH"]
assert "cool:python:path" in python_path
assert "paths" in python_path
def test_local_agent_ping_exception(monkeypatch, runner_token):
api = MagicMock()
api.ping.return_value = True
api.ping.side_effect = Exception()
monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient",
MagicMock(return_value=api))
with pytest.raises(Exception):
agent = LocalAgent()
def test_local_agent_config_options_populated(monkeypatch, runner_token):
api = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", api)
with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = LocalAgent(base_url="url", no_pull=True)
assert agent.client.get_auth_token() == "TEST_TOKEN"
assert agent.logger
assert agent.no_pull
assert api.call_args[1]["base_url"] == "url"
