def test_local_agent_init(runner_token):
agent = LocalAgent()
assert agent
assert set(agent.labels) == {
socket.gethostname(),
"azure-flow-storage",
"s3-flow-storage",
"gcs-flow-storage",
}
assert agent.name == "agent"
def test_local_agent_deploy_raises_unsupported_storage(monkeypatch):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent()
with pytest.raises(TypeError, match="Unsupported Storage type: Docker"):
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_populate_env_vars_uses_user_provided_env_vars_removes_nones(
runner_token):
with set_temporary_config({
"cloud.api": "api",
"logging.log_to_cloud": True,
"cloud.agent.auth_token": "token",
}):
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_start_max_polls_zero(monkeypatch, runner_token):
on_shutdown = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.LocalAgent.on_shutdown", on_shutdown)
agent_process = MagicMock()
monkeypatch.setattr("prefect.agent.agent.Agent.agent_process", agent_process)
agent_connect = MagicMock(return_value="id")
monkeypatch.setattr("prefect.agent.agent.Agent.agent_connect", agent_connect)
heartbeat = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.LocalAgent.heartbeat", heartbeat)
agent = LocalAgent(max_polls=0)
agent.start()
assert on_shutdown.call_count == 1
assert agent_process.call_count == 0
assert heartbeat.call_count == 0
def test_local_agent_config_options_hostname(runner_token):
with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = LocalAgent(labels=["test_label"])
assert set(agent.labels) == {
"test_label",
socket.gethostname(),
"azure-flow-storage",
"s3-flow-storage",
"gcs-flow-storage",
}
def test_local_agent_deploy_no_existing_python_path(monkeypatch, runner_token):
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": "id",
}
)
)
assert popen.called
assert len(agent.processes) == 1
assert "paths" in popen.call_args[1]["env"]["PYTHONPATH"]
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"
def test_local_agent_deploy_flows_storage_continues(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": Local().serialize()}), "id": "id"}
)
]
)
assert not api.pull.called
def test_local_agent_heartbeat(monkeypatch, returncode, show_flow_logs, logs):
popen = MockPopen()
# expect a process to be called with the following command (with specified behavior)
popen.set_command(
[sys.executable, "-m", "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 install(label, env, import_paths, **kwargs):
"""Generate a supervisord.conf file for a Local agent"""
from prefect.agent.local import LocalAgent
conf = LocalAgent.generate_supervisor_conf(
labels=sorted(set(label)),
env_vars=dict(e.split("=", 1) for e in env),
import_paths=list(import_paths),
**kwargs,
)
click.echo(conf)
def test_local_agent_config_options(monkeypatch, runner_token):
api = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", api)
monkeypatch.setattr("prefect.agent.local.agent.platform", "osx")
with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = LocalAgent()
assert agent.client.get_auth_token() == "TEST_TOKEN"
assert agent.logger
assert not agent.no_pull
assert api.call_args[1]["base_url"] == "unix://var/run/docker.sock"
def test_local_agent_daemon_url_responds_to_system(monkeypatch, runner_token):
api = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", api)
monkeypatch.setattr("prefect.agent.local.agent.platform", "win32")
with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = LocalAgent()
assert agent.client.get_auth_token() == "TEST_TOKEN"
assert agent.logger
assert not agent.no_pull
assert api.call_args[1]["base_url"] == "npipe:////./pipe/docker_engine"
def test_local_agent_deploy_null_or_univeral_run_config(
monkeypatch, run_config):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent()
agent.deploy_flow(flow_run=GraphQLResult(
{
"id": "id",
"flow": {
"storage": Local().serialize(),
"id": "foo",
"core_version": "0.13.0",
},
"run_config": run_config.serialize() if run_config else None,
}, ))
assert popen.called
assert len(agent.processes) == 1
def test_local_agent_deploy_storage_raises_not_supported_storage(
monkeypatch, runner_token):
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"
}
}, ))
assert not popen.called
assert len(agent.processes) == 0
def test_local_agent_deduplicates_labels(runner_token):
agent = LocalAgent(labels=["azure-flow-storage"])
assert agent
assert set(agent.labels) == {
socket.gethostname(),
"azure-flow-storage",
"s3-flow-storage",
"gcs-flow-storage",
"github-flow-storage",
}
assert len(agent.labels) == len(set(agent.labels))
def test_populate_env_vars(monkeypatch, backend, config_with_api_key):
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(TEST_FLOW_RUN_DATA)
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__API_KEY":
config_with_api_key.cloud.api_key,
"PREFECT__CLOUD__TENANT_ID":
config_with_api_key.cloud.tenant_id,
"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__CLOUD__SEND_FLOW_RUN_LOGS":
"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_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_init(cloud_api):
agent = LocalAgent()
assert agent
assert set(agent.labels) == {
socket.gethostname(),
"azure-flow-storage",
"s3-flow-storage",
"gcs-flow-storage",
"github-flow-storage",
"webhook-flow-storage",
}
assert agent.name == "agent"
def test_local_agent_no_pull(monkeypatch, runner_token):
api = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", api)
agent = LocalAgent()
assert not agent.no_pull
agent = LocalAgent(no_pull=True)
assert agent.no_pull
with context(no_pull=True):
agent = LocalAgent()
assert agent.no_pull
with context(no_pull=False):
agent = LocalAgent(no_pull=True)
assert agent.no_pull
with context(no_pull=False):
agent = LocalAgent(no_pull=False)
assert not agent.no_pull
def test_populate_env_vars_includes_agent_labels(cloud_api):
with set_temporary_config(
{
"cloud.api": "api",
"logging.log_to_cloud": True,
"cloud.agent.auth_token": "token",
}
):
agent = LocalAgent(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__AUTH_TOKEN": "token",
"PREFECT__CLOUD__AGENT__LABELS": str(
[
"42",
"marvin",
socket.gethostname(),
"azure-flow-storage",
"gcs-flow-storage",
"s3-flow-storage",
"github-flow-storage",
"webhook-flow-storage",
"gitlab-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": "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_local_agent_config_options(config_with_api_key):
agent = LocalAgent(
name="test",
labels=["test_label"],
import_paths=["test_path"],
)
assert agent.name == "test"
assert agent.client.api_key == config_with_api_key.cloud.api_key
assert agent.logger
assert agent.log_to_cloud is True
assert agent.processes == set()
assert agent.import_paths == ["test_path"]
assert set(agent.labels) == {"test_label", *DEFAULT_AGENT_LABELS}
def test_local_agent_config_options():
agent = LocalAgent(
name="test",
labels=["test_label"],
import_paths=["test_path"],
)
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) == {"test_label", *DEFAULT_AGENT_LABELS}
def test_local_agent_deploy_processes_gitlab_storage(monkeypatch, cloud_api):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent()
gitlab = GitLab("test/repo", path="path/to/flow.py")
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": gitlab.serialize(),
"id": "foo",
}
),
"id": "id",
}
)
)
assert popen.called
assert len(agent.processes) == 1
def test_local_agent_no_pull(monkeypatch):
api = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", api)
with set_temporary_config({"cloud.agent.auth_token": "token"}):
agent = LocalAgent()
assert not agent.no_pull
agent = LocalAgent(no_pull=True)
assert agent.no_pull
with context(no_pull=True):
agent = LocalAgent()
assert agent.no_pull
with context(no_pull=False):
agent = LocalAgent(no_pull=True)
assert agent.no_pull
with context(no_pull=False):
agent = LocalAgent(no_pull=False)
assert not agent.no_pull
def test_local_agent_deploy_processes_valid_storage(storage, monkeypatch):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent()
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": storage.serialize(),
"id": "foo",
"core_version": "0.13.0",
}
),
"id": "id",
}
)
)
assert popen.called
assert len(agent.processes) == 1
def test_local_agent_config_no_storage_labels(cloud_api):
with set_temporary_config({
"cloud.agent.auth_token": "TEST_TOKEN",
"logging.log_to_cloud": True
}):
agent = LocalAgent(
labels=["test_label"],
storage_labels=False,
)
assert set(agent.labels) == {
socket.gethostname(),
"test_label",
}
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_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_import_paths(monkeypatch, cloud_api):
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"]
