def test_populate_env_vars_includes_agent_labels(api):
agent = DockerAgent(labels=["42", "marvin"])
env_vars = agent.populate_env_vars(
GraphQLResult({"id": "id", "name": "name", "flow": {"id": "foo"}})
)
assert env_vars["PREFECT__CLOUD__AGENT__LABELS"] == "['42', 'marvin']"
def test_environment_has_api_key_from_disk(api, monkeypatch, tenant_id):
"""Check that the API key is passed through from the on disk cache"""
monkeypatch.setattr(
"prefect.Client.load_auth_from_disk",
MagicMock(return_value={
"api_key": "TEST_KEY",
"tenant_id": tenant_id
}),
)
agent = DockerAgent()
agent.client._get_auth_tenant = MagicMock(return_value="ID")
env = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
}
}),
"test-image",
)
assert env["PREFECT__CLOUD__API_KEY"] == "TEST_KEY"
assert env["PREFECT__CLOUD__AUTH_TOKEN"] == "TEST_KEY"
assert env["PREFECT__CLOUD__TENANT_ID"] == "ID"
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__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_sets_log_to_cloud(flag, 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_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"}})
)
assert env_vars["AUTH_THING"] == "foo"
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_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_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_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_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_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_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_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_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_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_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_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"
