From 216ae61d621a4609ba3ea2ec1838d7aac5d28daf Mon Sep 17 00:00:00 2001 From: Zain Patel <52913697+ZainPatelQB@users.noreply.github.com> Date: Mon, 24 Feb 2020 19:08:38 +0000 Subject: [PATCH] Refactor KEDRO_ENV (#447) --- kedro/context/__init__.py | 7 +-- kedro/context/context.py | 6 ++- .../profile_default/startup/00-kedro-init.py | 6 +-- .../{{ cookiecutter.repo_name }}/kedro_cli.py | 9 ++-- tests/template/test_kedro_cli.py | 53 ++++--------------- tests/template/test_load_context.py | 10 ++++ 6 files changed, 30 insertions(+), 61 deletions(-) diff --git a/kedro/context/__init__.py b/kedro/context/__init__.py index 169f509867..eafd0bffac 100644 --- a/kedro/context/__init__.py +++ b/kedro/context/__init__.py @@ -30,9 +30,4 @@ project context. """ -from .context import ( # NOQA - KEDRO_ENV_VAR, - KedroContext, - KedroContextError, - load_context, -) +from .context import KedroContext, KedroContextError, load_context # NOQA diff --git a/kedro/context/context.py b/kedro/context/context.py index 4834bcf155..953cbe17b6 100644 --- a/kedro/context/context.py +++ b/kedro/context/context.py @@ -48,8 +48,6 @@ from kedro.utils import load_obj from kedro.versioning import Journal -KEDRO_ENV_VAR = "KEDRO_ENV" - def _version_mismatch_error(context_version) -> str: return ( @@ -562,6 +560,10 @@ def load_context(project_path: Union[str, Path], **kwargs) -> KedroContext: ) os.chdir(str(project_path)) # Move to project root + # update kwargs with env from the environment variable (defaults to None if not set) + # need to do this because some CLI command (e.g `kedro run`) defaults to passing in `env=None` + kwargs["env"] = kwargs.get("env") or os.getenv("KEDRO_ENV") + # Instantiate the context after changing the cwd for logging to be properly configured. context = context_class(project_path, **kwargs) return context diff --git a/kedro/template/{{ cookiecutter.repo_name }}/.ipython/profile_default/startup/00-kedro-init.py b/kedro/template/{{ cookiecutter.repo_name }}/.ipython/profile_default/startup/00-kedro-init.py index 259e5e5154..f0943c4190 100644 --- a/kedro/template/{{ cookiecutter.repo_name }}/.ipython/profile_default/startup/00-kedro-init.py +++ b/kedro/template/{{ cookiecutter.repo_name }}/.ipython/profile_default/startup/00-kedro-init.py @@ -1,5 +1,4 @@ import logging.config -import os import sys from pathlib import Path @@ -19,7 +18,7 @@ def reload_kedro(path, line=None): try: import kedro.config.default_logger - from kedro.context import KEDRO_ENV_VAR, load_context + from kedro.context import load_context from kedro.cli.jupyter import collect_line_magic except ImportError: logging.error( @@ -31,8 +30,7 @@ def reload_kedro(path, line=None): try: path = path or project_path logging.debug("Loading the context from %s", str(path)) - - context = load_context(path, env=os.getenv(KEDRO_ENV_VAR)) + context = load_context(path) catalog = context.catalog # remove cached user modules diff --git a/kedro/template/{{ cookiecutter.repo_name }}/kedro_cli.py b/kedro/template/{{ cookiecutter.repo_name }}/kedro_cli.py index d88abc0fef..9103d87a0e 100755 --- a/kedro/template/{{ cookiecutter.repo_name }}/kedro_cli.py +++ b/kedro/template/{{ cookiecutter.repo_name }}/kedro_cli.py @@ -51,7 +51,7 @@ forward_command, python_call, ) -from kedro.context import KEDRO_ENV_VAR, load_context +from kedro.context import load_context from kedro.runner import SequentialRunner from kedro.utils import load_obj @@ -216,7 +216,6 @@ def cli(): type=str, default=None, multiple=False, - envvar=KEDRO_ENV_VAR, help=ENV_ARG_HELP, ) @click.option("--tag", "-t", type=str, multiple=True, help=TAG_ARG_HELP) @@ -456,12 +455,12 @@ def _build_jupyter_command( def _build_jupyter_env(kedro_env: str) -> Dict[str, Any]: """Build the environment dictionary that gets injected into the subprocess running Jupyter. Since the subprocess has access only to the environment variables passed - in, we need to copy the current environment and add ``KEDRO_ENV_VAR``. + in, we need to copy the current environment and add ``KEDRO_ENV``. """ if not kedro_env: return {} jupyter_env = os.environ.copy() - jupyter_env[KEDRO_ENV_VAR] = kedro_env + jupyter_env["KEDRO_ENV"] = kedro_env return {"env": jupyter_env} @@ -484,7 +483,6 @@ def jupyter(): type=str, default=None, multiple=False, - envvar=KEDRO_ENV_VAR, help=ENV_ARG_HELP, ) def jupyter_notebook(ip, all_kernels, env, idle_timeout, args): @@ -512,7 +510,6 @@ def jupyter_notebook(ip, all_kernels, env, idle_timeout, args): type=str, default=None, multiple=False, - envvar=KEDRO_ENV_VAR, help=ENV_ARG_HELP, ) def jupyter_lab(ip, all_kernels, env, idle_timeout, args): diff --git a/tests/template/test_kedro_cli.py b/tests/template/test_kedro_cli.py index 1cd63300e6..783fce9ada 100644 --- a/tests/template/test_kedro_cli.py +++ b/tests/template/test_kedro_cli.py @@ -37,7 +37,6 @@ import pytest from click.testing import CliRunner -from kedro.context import KEDRO_ENV_VAR from kedro.runner import ParallelRunner, SequentialRunner @@ -238,7 +237,10 @@ def test_run_with_config( ({}, {}), ({"params": {"foo": "baz"}}, {"foo": "baz"}), ({"params": "foo:baz"}, {"foo": "baz"}), - ({"params": {"foo": "123.45", "baz": "678", "bar": 9}}, {"foo": "123.45", "baz": "678", "bar": 9}), + ( + {"params": {"foo": "123.45", "baz": "678", "bar": 9}}, + {"foo": "123.45", "baz": "678", "bar": 9}, + ), ], indirect=["fake_run_config_with_params"], ) @@ -268,17 +270,6 @@ def test_run_with_params_in_config( Path.cwd(), env=mocker.ANY, extra_params=expected ) - def test_run_env_environment_var( - self, fake_kedro_cli, fake_load_context, fake_repo_path, monkeypatch, mocker - ): - monkeypatch.setenv("KEDRO_ENV", "my_special_env") - result = CliRunner().invoke(fake_kedro_cli.cli, ["run"]) - assert not result.exit_code - - fake_load_context.assert_called_once_with( - Path.cwd(), env="my_special_env", extra_params=mocker.ANY - ) - @pytest.mark.parametrize( "cli_arg,expected_extra_params", [ @@ -614,6 +605,7 @@ def test_help(self, help_flag, fake_kedro_cli, fake_ipython_message): def test_env( self, env_flag, fake_kedro_cli, python_call_mock, default_jupyter_options ): + """This tests passing an environment variable to the jupyter subprocess.""" result = CliRunner().invoke( fake_kedro_cli.cli, ["jupyter", "notebook", env_flag, "my_special_env"] ) @@ -622,21 +614,8 @@ def test_env( args, kwargs = python_call_mock.call_args assert args == default_jupyter_options assert "env" in kwargs - assert KEDRO_ENV_VAR in kwargs["env"] - assert kwargs["env"][KEDRO_ENV_VAR] == "my_special_env" - - def test_env_environment_variable( - self, fake_kedro_cli, python_call_mock, monkeypatch, default_jupyter_options - ): - monkeypatch.setenv("KEDRO_ENV", "my_special_env") - result = CliRunner().invoke(fake_kedro_cli.cli, ["jupyter", "notebook"]) - assert not result.exit_code - - args, kwargs = python_call_mock.call_args - assert args == default_jupyter_options - assert "env" in kwargs - assert KEDRO_ENV_VAR in kwargs["env"] - assert kwargs["env"][KEDRO_ENV_VAR] == "my_special_env" + assert "KEDRO_ENV" in kwargs["env"] + assert kwargs["env"]["KEDRO_ENV"] == "my_special_env" class TestJupyterLabCommand: @@ -707,6 +686,7 @@ def test_help(self, help_flag, fake_kedro_cli, fake_ipython_message): def test_env( self, env_flag, fake_kedro_cli, python_call_mock, default_jupyter_options ): + """This tests passing an environment variable to the jupyter subprocess.""" result = CliRunner().invoke( fake_kedro_cli.cli, ["jupyter", "lab", env_flag, "my_special_env"] ) @@ -715,21 +695,8 @@ def test_env( args, kwargs = python_call_mock.call_args assert args == default_jupyter_options assert "env" in kwargs - assert KEDRO_ENV_VAR in kwargs["env"] - assert kwargs["env"][KEDRO_ENV_VAR] == "my_special_env" - - def test_env_environment_variable( - self, fake_kedro_cli, python_call_mock, monkeypatch, default_jupyter_options - ): - monkeypatch.setenv("KEDRO_ENV", "my_special_env") - result = CliRunner().invoke(fake_kedro_cli.cli, ["jupyter", "lab"]) - assert not result.exit_code - - args, kwargs = python_call_mock.call_args - assert args == default_jupyter_options - assert "env" in kwargs - assert KEDRO_ENV_VAR in kwargs["env"] - assert kwargs["env"][KEDRO_ENV_VAR] == "my_special_env" + assert "KEDRO_ENV" in kwargs["env"] + assert kwargs["env"]["KEDRO_ENV"] == "my_special_env" class TestConvertNotebookCommand: diff --git a/tests/template/test_load_context.py b/tests/template/test_load_context.py index b4b19b26bb..ad63bfad7d 100644 --- a/tests/template/test_load_context.py +++ b/tests/template/test_load_context.py @@ -47,6 +47,16 @@ def test_valid_context(self, mocker, fake_repo_path): assert str(fake_repo_path.resolve() / "src") in sys.path assert os.getcwd() == str(fake_repo_path.resolve()) + def test_valid_context_with_env(self, mocker, monkeypatch, fake_repo_path): + """Test getting project context when Kedro config environment is specified in the environment variable.""" + # Disable logging.config.dictConfig in KedroContext._setup_logging as + # it changes logging.config and affects other unit tests + mocker.patch("logging.config.dictConfig") + mocker.patch("kedro.config.config.ConfigLoader.get") + monkeypatch.setenv("KEDRO_ENV", "my_fake_env") + result = load_context(str(fake_repo_path)) + assert result.env == "my_fake_env" + def test_invalid_path(self, tmp_path): """Test for loading context from an invalid path. """ other_path = tmp_path / "other"