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(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_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_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(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_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_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_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_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_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_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_environment_has_tenant_id_from_server(config_with_api_key):
tenant_id = uuid.uuid4()
with set_temporary_config({"cloud.tenant_id": None}):
agent = LocalAgent()
agent.client._get_auth_tenant = MagicMock(return_value=tenant_id)
env = agent.populate_env_vars(TEST_FLOW_RUN_DATA)
assert env["PREFECT__CLOUD__API_KEY"] == config_with_api_key.cloud.api_key
assert env["PREFECT__CLOUD__TENANT_ID"] == tenant_id
def test_environment_has_api_key_from_config(config_with_api_key):
"""Check that the API key is passed through from the config via environ"""
agent = LocalAgent()
agent.client._get_auth_tenant = MagicMock(return_value="ID")
env = agent.populate_env_vars(TEST_FLOW_RUN_DATA)
assert env["PREFECT__CLOUD__API_KEY"] == config_with_api_key.cloud.api_key
assert env[
"PREFECT__CLOUD__TENANT_ID"] == config_with_api_key.cloud.tenant_id
def test_populate_env_vars_uses_user_provided_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=dict(AUTH_THING="foo"))
env_vars = agent.populate_env_vars(GraphQLResult({"id": "id"}))
assert env_vars["AUTH_THING"] == "foo"
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(
GraphQLResult({"id": "id", "name": "name", "flow": {"id": "foo"}})
)
python_path = env_vars["PYTHONPATH"]
assert "cool:python:path" in python_path
assert "paths" in python_path
def test_populate_env_vars_import_paths():
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_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"}))
assert "MISSING_VAR" not in env_vars
def test_environment_has_api_key_from_config(tenant_id):
"""Check that the API key is passed through from the config via environ"""
with set_temporary_config({
"cloud.api_key": "TEST_KEY",
"cloud.tenant_id": tenant_id
}):
agent = LocalAgent()
env = agent.populate_env_vars(TEST_FLOW_RUN_DATA)
assert env["PREFECT__CLOUD__API_KEY"] == "TEST_KEY"
assert env.get("PREFECT__CLOUD__TENANT_ID") == tenant_id
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_environment_has_api_key_from_config(tenant_id):
"""Check that the API key is passed through from the config via environ"""
with set_temporary_config({
"cloud.api_key": "TEST_KEY",
"cloud.tenant_id": tenant_id,
"cloud.agent.auth_token": None,
}):
agent = LocalAgent()
agent.client._get_auth_tenant = MagicMock(return_value="ID")
env = agent.populate_env_vars(TEST_FLOW_RUN_DATA)
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_environment_has_api_key_from_disk(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 = LocalAgent()
env = agent.populate_env_vars(TEST_FLOW_RUN_DATA)
assert env["PREFECT__CLOUD__API_KEY"] == "TEST_KEY"
assert env.get("PREFECT__CLOUD__TENANT_ID") == tenant_id
def test_environment_has_api_key_from_disk(monkeypatch):
"""Check that the API key is passed through from the on disk cache"""
tenant_id = str(uuid.uuid4())
monkeypatch.setattr(
"prefect.Client.load_auth_from_disk",
MagicMock(return_value={
"api_key": "TEST_KEY",
"tenant_id": tenant_id
}),
)
with set_temporary_config({"cloud.tenant_id": None}):
agent = LocalAgent()
env = agent.populate_env_vars(TEST_FLOW_RUN_DATA)
assert env["PREFECT__CLOUD__API_KEY"] == "TEST_KEY"
assert env["PREFECT__CLOUD__TENANT_ID"] == tenant_id
def test_prefect_logging_level_override_logic(config, agent_env_vars,
run_config_env_vars,
expected_logging_level, tmpdir):
with set_temporary_config(config):
agent = LocalAgent(env_vars=agent_env_vars)
run = LocalRun(working_dir=str(tmpdir), env=run_config_env_vars)
env_vars = agent.populate_env_vars(
GraphQLResult({
"id": "id",
"name": "name",
"flow": {
"id": "foo"
},
"run_config": run.serialize(),
}),
run,
)
assert env_vars["PREFECT__LOGGING__LEVEL"] == expected_logging_level
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_no_existing_python_path(monkeypatch):
monkeypatch.delenv("PYTHONPATH", raising=False)
agent = LocalAgent(import_paths=["paths"])
env_vars = agent.populate_env_vars(TEST_FLOW_RUN_DATA)
assert "paths" in env_vars["PYTHONPATH"]
def test_populate_env_vars_import_paths():
agent = LocalAgent(import_paths=["paths"])
env_vars = agent.populate_env_vars(TEST_FLOW_RUN_DATA)
assert "paths" in env_vars["PYTHONPATH"]
def test_populate_env_vars_includes_agent_labels():
agent = LocalAgent(labels=["42", "marvin"])
env_vars = agent.populate_env_vars(TEST_FLOW_RUN_DATA)
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(TEST_FLOW_RUN_DATA)
assert "MISSING_VAR" not in env_vars
def test_populate_env_vars_from_agent_config():
agent = LocalAgent(env_vars=dict(AUTH_THING="foo"))
env_vars = agent.populate_env_vars(TEST_FLOW_RUN_DATA)
assert env_vars["AUTH_THING"] == "foo"
