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"