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_deploy_flow_reg_allow_list_not_allowed(monkeypatch, cloud_api):
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(reg_allow_list=["test1"])
with pytest.raises(ValueError) as error:
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"id": "foo",
"storage": Docker(
registry_url="test2", image_name="name", image_tag="tag"
).serialize(),
"environment": LocalEnvironment().serialize(),
"core_version": "0.13.0",
}
),
"id": "id",
"name": "name",
}
)
)
expected_error = (
"Trying to pull image from a Docker registry 'test2'"
" which is not in the reg_allow_list"
)
assert not api.pull.called
assert not api.create_container.called
assert not api.start.called
assert str(error.value) == expected_error
def test_docker_agent_deploy_flow_show_flow_logs(api, monkeypatch):
process = MagicMock()
monkeypatch.setattr("multiprocessing.Process", process)
agent = DockerAgent(show_flow_logs=True)
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",
}))
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_deploy_flow(core_version, command, 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":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
LocalEnvironment().serialize(),
"core_version":
core_version,
}),
"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"] == command
assert api.create_container.call_args[1]["host_config"][
"AutoRemove"] is True
assert api.start.call_args[1]["container"] == "container_id"
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_networks_as_modes_can_be_overriden_by_run_config(
api, network):
api.create_networking_config.return_value = {network: "config1"}
api.create_endpoint_config.return_value = "endpoint-config"
agent = DockerAgent(networks=[network])
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(),
"core_version":
"0.13.0",
}),
"run_config":
DockerRun(host_config={
"network_mode": "foobar"
}).serialize(),
"id":
"id",
"name":
"name",
}))
assert network in agent.networks
api.create_networking_config.assert_called_once_with(
{network: "endpoint-config"})
_, container_create_kwargs = api.create_container.call_args
assert container_create_kwargs["networking_config"] == {network: "config1"}
_, host_config_kwargs = api.create_host_config.call_args
assert host_config_kwargs["network_mode"] == "foobar"
def test_environment_has_tenant_id_from_server(api, config_with_api_key):
tenant_id = uuid.uuid4()
with set_temporary_config({"cloud.tenant_id": None}):
agent = DockerAgent()
agent.client._get_auth_tenant = MagicMock(return_value=tenant_id)
env = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}),
"test-image",
)
assert env["PREFECT__CLOUD__API_KEY"] == config_with_api_key.cloud.api_key
assert env[
"PREFECT__CLOUD__AUTH_TOKEN"] == config_with_api_key.cloud.api_key
assert env["PREFECT__CLOUD__TENANT_ID"] == tenant_id
def test_docker_agent_start_max_polls(max_polls, api, monkeypatch, runner_token):
on_shutdown = MagicMock()
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent.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.docker.agent.DockerAgent.heartbeat", heartbeat)
agent = DockerAgent(max_polls=max_polls)
agent.start()
assert agent_connect.call_count == 1
assert agent_process.call_count == max_polls
assert heartbeat.call_count == 1
assert on_shutdown.call_count == 1
def test_populate_env_vars(api):
agent = DockerAgent()
env_vars = agent.populate_env_vars(
GraphQLResult({"id": "id", "name": "name", "flow": {"id": "foo"}}), "test-image"
)
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__CONTEXT__IMAGE": "test-image",
"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_unsupported_run_config(api):
agent = DockerAgent()
with pytest.raises(TypeError,
match="Unsupported RunConfig type: LocalRun"):
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage": Local().serialize(),
"run_config": LocalRun().serialize(),
"id": "foo",
"core_version": "0.13.0",
}),
"id":
"id",
"name":
"name",
"version":
"version",
}))
assert not api.pull.called
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(),
}
),
"test-image",
run_config=run,
)
assert env_vars["KEY1"] == "VAL1"
assert env_vars["KEY2"] == "OVERRIDE"
assert env_vars["PREFECT__LOGGING__LEVEL"] == "TEST"
def test_docker_agent_networks(api):
api.create_networking_config.return_value = {"test-network-1": "config1"}
api.create_endpoint_config.return_value = "endpoint-config"
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(),
"core_version":
"0.13.0",
}),
"id":
"id",
"name":
"name",
}))
assert "test-network-1" in agent.networks
assert "test-network-2" in agent.networks
api.create_networking_config.assert_called_once_with(
{"test-network-1": "endpoint-config"})
_, kwargs = api.create_container.call_args
assert kwargs["networking_config"] == {"test-network-1": "config1"}
api.connect_container_to_network.assert_called_once_with(
container=ANY, net_id="test-network-2")
_, host_config_kwargs = api.create_host_config.call_args
assert "network_mode" not in host_config_kwargs
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_docker_agent_deploy_flow(core_version, command, api):
agent = DockerAgent()
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":
core_version,
}),
"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"] == command
assert api.create_container.call_args[1]["host_config"][
"AutoRemove"] is True
assert api.create_container.call_args[1]["labels"] == {
"io.prefect.flow-id": "foo",
"io.prefect.flow-name": "flow-name",
"io.prefect.flow-run-id": "id",
}
assert api.start.call_args[1]["container"] == "container_id"
def test_environment_has_api_key_from_config(api, tenant_id):
with set_temporary_config({
"cloud.api_key": "TEST_KEY",
"cloud.tenant_id": tenant_id,
"cloud.agent.auth_token": None,
}):
agent = DockerAgent()
agent.client._get_auth_tenant = MagicMock(return_value="ID")
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}),
"test-image",
)
assert env_vars["PREFECT__CLOUD__API_KEY"] == "TEST_KEY"
assert env_vars["PREFECT__CLOUD__AUTH_TOKEN"] == "TEST_KEY"
assert env_vars["PREFECT__CLOUD__TENANT_ID"] == "ID"
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_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",
"environment": LocalEnvironment().serialize(),
"core_version": "0.13.0",
}),
"id":
"id",
"name":
"name",
"version":
"version",
}))
assert not api.pull.called
def test_docker_agent_deploy_flow_sets_container_name_with_slugify(
api, run_name, container_name):
"""
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
"""
agent = DockerAgent()
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id": "foo",
"name": "flow-name",
"storage": Local().serialize(),
"core_version": "0.13.0",
}),
"id":
"id",
"name":
run_name,
}))
assert api.create_container.call_args[1]["name"] == container_name
def test_docker_agent_deploy_with_no_interface_check_linux(
monkeypatch, cloud_api, linux_platform
):
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),
)
get_ip = MagicMock()
monkeypatch.setattr("prefect.agent.docker.agent.get_docker_ip", get_ip)
agent = DockerAgent(docker_interface=False)
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"id": "foo",
"storage": Docker(
registry_url="", image_name="name", image_tag="tag"
).serialize(),
"environment": LocalEnvironment().serialize(),
"core_version": "0.13.0",
}
),
"id": "id",
"name": "name",
}
)
)
assert not get_ip.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_docker_agent_deploy_flow_no_pull_using_environment_metadata(api):
agent = DockerAgent(no_pull=True)
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"id": "foo",
"storage": Local().serialize(),
"environment": LocalEnvironment(
metadata={"image": "name: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_api_url_uses_server_network(api, backend):
"""
If the `prefect-server` network is provided and the backend is 'server' then we
will replace the API url with 'apollo' instead of 'host.docker.internal' to make
use of the network for connections to the API
"""
agent = DockerAgent(networks=["prefect-server"])
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
},
"run_config": {},
}),
"test-image",
)
if backend == "server":
assert env_vars["PREFECT__CLOUD__API"] == "http://apollo:4200"
else:
assert env_vars["PREFECT__CLOUD__API"] == "https://api.prefect.io"
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_deploy_flow_no_registry_does_not_pull(api):
agent = DockerAgent()
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"id": "foo",
"storage": Docker(
registry_url="", image_name="name", image_tag="tag"
).serialize(),
"environment": LocalEnvironment().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_ping(api):
DockerAgent()
assert api.ping.called
def test_docker_agent_ping_exception(api):
api.ping.side_effect = Exception()
with pytest.raises(Exception):
DockerAgent()
def test_docker_agent_is_named_volume_win32(monkeypatch, api, path, result):
monkeypatch.setattr("prefect.agent.docker.agent.platform", "win32")
agent = DockerAgent()
assert agent._is_named_volume_win32(path) == result
def test_docker_agent_init_volume_empty_options(api):
agent = DockerAgent()
assert agent
assert agent.named_volumes == []
assert agent.container_mount_paths == []
assert agent.host_spec == {}
def test_docker_agent_heartbeat_gocase(api):
agent = DockerAgent()
agent.heartbeat()
assert api.ping.call_count == 2
def test_docker_agent_init(api):
agent = DockerAgent()
assert agent
assert agent.agent_config_id is None
assert agent.labels == []
assert agent.name == "agent"
