def test_docker_agent_parse_volume_spec_win(api, candidate, named_volumes,
container_mount_paths, host_spec):
agent = DockerAgent()
(
actual_named_volumes,
actual_container_mount_paths,
actual_host_spec,
) = agent._parse_volume_spec_win32(candidate)
assert actual_named_volumes == named_volumes
assert actual_container_mount_paths == container_mount_paths
assert actual_host_spec == host_spec
def test_environment_has_agent_token_from_config(api, config_with_token):
agent = DockerAgent()
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}), "test-image")
assert env_vars["PREFECT__CLOUD__AUTH_TOKEN"] == "TEST_TOKEN"
def test_populate_env_vars_from_agent_config(api):
agent = DockerAgent(env_vars=dict(AUTH_THING="foo"))
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}), "test-image")
assert env_vars["AUTH_THING"] == "foo"
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":
"",
"PREFECT__CLOUD__API_KEY":
"",
"PREFECT__CLOUD__TENANT_ID":
"",
"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_api_url_can_be_overridden_at_agent_level(api):
agent = DockerAgent(env_vars={"PREFECT__CLOUD__API": "FOO"})
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
},
}),
"test-image",
)
assert env_vars["PREFECT__CLOUD__API"] == "FOO"
def test_docker_agent_deploy_flow_sets_container_name_with_index(
api, collision_count):
"""
Asserts that the container name is set to the flow run name and that collisions with
existing containers with the same name is handled by adding an index
"""
if collision_count:
# Add the basic name first
existing_names = ["flow-run-name"]
for i in range(1, collision_count):
existing_names.append(f"flow-run-name-{i}")
else:
existing_names = []
def fail_if_name_exists(*args, **kwargs):
if kwargs.get("name") in existing_names:
raise docker.errors.APIError(
"Conflict. The container name 'foobar' is already in use")
return {}
api.create_container = MagicMock(side_effect=fail_if_name_exists)
agent = DockerAgent()
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"name":
"flow-name",
"storage":
Local().serialize(),
"environment":
LocalEnvironment(metadata={
"image": "repo/name:tag"
}).serialize(),
"core_version":
"0.13.0",
}),
"id":
"id",
"name":
"flow-run-name",
}))
expected_name = ("flow-run-name" if not collision_count else
f"flow-run-name-{collision_count}")
assert api.create_container.call_args[1]["name"] == expected_name
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_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_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_docker_agent_deploy_flow_show_flow_logs(monkeypatch, cloud_api):
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(),
"core_version":
"0.13.0",
}),
"id":
"id",
"name":
"name",
}))
process_kwargs = dict(
target=_stream_container_logs,
kwargs={
"base_url": agent.base_url,
"container_id": "container_id"
},
)
process.assert_called_with(**process_kwargs)
# Check all arguments to `multiprocessing.Process` are pickleable
assert pickle.loads(pickle.dumps(process_kwargs)) == process_kwargs
assert len(agent.processes) == 1
assert api.create_container.called
assert api.start.called
def test_docker_agent_shutdown_terminates_child_processes(monkeypatch, cloud_api):
monkeypatch.setattr("prefect.agent.agent.Client", MagicMock())
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),
)
proc = MagicMock(is_alive=MagicMock(return_value=True))
agent = DockerAgent(show_flow_logs=True)
agent.processes = [proc]
agent.on_shutdown()
assert proc.is_alive.called
assert proc.terminate.called
def test_api_url_can_be_overridden_with_run_config(api):
agent = DockerAgent(env_vars={"PREFECT__CLOUD__API": "FOO"})
run = DockerRun(env={"PREFECT__CLOUD__API": "BAR"}, )
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
},
"run_config": run.serialize(),
}),
"test-image",
run_config=run,
)
assert env_vars["PREFECT__CLOUD__API"] == "BAR"
def test_docker_agent_no_pull(api):
agent = DockerAgent()
assert not agent.no_pull
agent = DockerAgent(no_pull=True)
assert agent.no_pull
with context(no_pull=True):
agent = DockerAgent()
assert agent.no_pull
with context(no_pull=False):
agent = DockerAgent(no_pull=True)
assert agent.no_pull
with context(no_pull=False):
agent = DockerAgent(no_pull=False)
assert not agent.no_pull
def test_docker_agent_config_options_populated(monkeypatch):
api = MagicMock()
monkeypatch.setattr("docker.APIClient", api)
agent = DockerAgent(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_docker_agent_ping(monkeypatch, cloud_api):
api = MagicMock()
api.ping.return_value = True
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
agent = DockerAgent()
assert api.ping.called
def test_docker_agent_heartbeat_exits_on_failure(monkeypatch, cloud_api, caplog):
api = MagicMock()
api.ping.return_value = True
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
agent = DockerAgent()
api.ping.return_value = False
agent.heartbeat()
agent.heartbeat()
agent.heartbeat()
agent.heartbeat()
agent.heartbeat()
with pytest.raises(SystemExit):
agent.heartbeat()
assert "Cannot reconnect to Docker daemon. Agent is shutting down." in caplog.text
assert api.ping.call_count == 7
def test_docker_agent_deploy_flow_uses_environment_metadata(
monkeypatch, cloud_api):
api = MagicMock()
api.ping.return_value = True
api.create_container.return_value = {"Id": "container_id"}
api.create_host_config.return_value = {"AutoRemove": True}
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
agent = DockerAgent()
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"storage":
Local().serialize(),
"environment":
LocalEnvironment(metadata={
"image": "repo/name:tag"
}).serialize(),
"core_version":
"0.13.0",
}),
"id":
"id",
"name":
"name",
}))
assert api.pull.called
assert api.create_container.called
assert api.start.called
assert api.create_host_config.call_args[1]["auto_remove"] is True
assert api.create_container.call_args[1][
"command"] == "prefect execute flow-run"
assert api.create_container.call_args[1]["host_config"][
"AutoRemove"] is True
assert api.start.call_args[1]["container"] == "container_id"
def test_environment_has_api_key_from_config(api, config_with_api_key):
agent = DockerAgent()
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}),
"test-image",
)
assert env_vars[
"PREFECT__CLOUD__API_KEY"] == config_with_api_key.cloud.api_key
assert env_vars[
"PREFECT__CLOUD__AUTH_TOKEN"] == config_with_api_key.cloud.api_key
assert env_vars[
"PREFECT__CLOUD__TENANT_ID"] == config_with_api_key.cloud.tenant_id
def test_docker_agent_ping_exception(monkeypatch, cloud_api):
api = MagicMock()
api.ping.return_value = True
api.ping.side_effect = Exception()
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
with pytest.raises(Exception):
agent = DockerAgent()
def test_docker_agent_init(monkeypatch, cloud_api):
api = MagicMock()
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
agent = DockerAgent()
assert agent
assert agent.labels == []
assert agent.name == "agent"
def test_docker_agent_config_options(platform, url, monkeypatch):
api = MagicMock()
monkeypatch.setattr("docker.APIClient", api)
monkeypatch.setattr("prefect.agent.docker.agent.platform", platform)
agent = DockerAgent(name="test")
assert agent.name == "test"
assert agent.client.get_auth_token() == "TEST_TOKEN"
assert agent.logger
assert not agent.no_pull
assert api.call_args[1]["base_url"] == url
def test_docker_agent_parse_volume_spec_win(
monkeypatch, cloud_api, candidate, named_volumes, container_mount_paths, host_spec
):
api = MagicMock()
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
agent = DockerAgent()
(
actual_named_volumes,
actual_container_mount_paths,
actual_host_spec,
) = agent._parse_volume_spec_win32(candidate)
assert actual_named_volumes == named_volumes
assert actual_container_mount_paths == container_mount_paths
assert actual_host_spec == host_spec
def test_populate_env_vars_from_run_config(api):
agent = DockerAgent(env_vars={"KEY1": "VAL1", "KEY2": "VAL2"})
run = DockerRun(
env={"KEY2": "OVERRIDE", "PREFECT__LOGGING__LEVEL": "TEST"},
)
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"
def test_docker_agent_deploy_flow_run_config(api, run_kind, has_docker_storage):
if has_docker_storage:
storage = Docker(
registry_url="testing", image_name="on-storage", image_tag="tag"
)
image = "testing/on-storage:tag"
else:
storage = Local()
image = "on-run-config" if run_kind == "docker" else "prefecthq/prefect:0.13.11"
if run_kind == "docker":
env = {"TESTING": "VALUE"}
run = DockerRun(image=image, env=env)
else:
env = {}
run = None if run_kind == "missing" else UniversalRun()
agent = DockerAgent()
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"id": "foo",
"name": "flow-name",
"storage": storage.serialize(),
"core_version": "0.13.11",
}
),
"run_config": run.serialize() if run else None,
"id": "id",
"name": "name",
}
)
)
assert api.create_container.called
assert api.create_container.call_args[0][0] == image
res_env = api.create_container.call_args[1]["environment"]
for k, v in env.items():
assert res_env[k] == v
def test_docker_agent_config_options(platform, url, monkeypatch,
config_with_api_key):
api = MagicMock()
monkeypatch.setattr("docker.APIClient", api)
monkeypatch.setattr("prefect.agent.docker.agent.platform", platform)
agent = DockerAgent(name="test")
assert agent.name == "test"
assert agent.client.api_key == config_with_api_key.cloud.api_key
assert agent.logger
assert not agent.no_pull
assert api.call_args[1]["base_url"] == url
def test_populate_env_vars(api):
agent = DockerAgent()
env_vars = agent.populate_env_vars(
GraphQLResult({"id": "id", "name": "name", "flow": {"id": "foo"}})
)
expected_vars = {
"PREFECT__CLOUD__API": "https://api.prefect.io",
"PREFECT__CLOUD__AUTH_TOKEN": "TEST_TOKEN",
"PREFECT__CLOUD__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_vars
def test_populate_env_vars_uses_user_provided_env_vars(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(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_docker_agent_init_volume_empty_options(monkeypatch, cloud_api):
api = MagicMock()
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
agent = DockerAgent()
assert agent
assert agent.named_volumes == []
assert agent.container_mount_paths == []
assert agent.host_spec == {}
def test_prefect_logging_level_override_logic(config, agent_env_vars,
run_config_env_vars,
expected_logging_level, api):
with set_temporary_config(config):
agent = DockerAgent(env_vars=agent_env_vars)
run = DockerRun(env=run_config_env_vars)
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
},
"run_config": run.serialize(),
}),
"test-image",
run_config=run,
)
assert env_vars["PREFECT__LOGGING__LEVEL"] == expected_logging_level
def test_docker_agent_config_options_populated(monkeypatch, cloud_api):
import docker # DockerAgent imports docker within the constructor
api = MagicMock()
monkeypatch.setattr("docker.APIClient", api)
with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = DockerAgent(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"
