From 71ea12877d81f1f4606ed3d97a75078d55a7062d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Peter=20Rowlands=20=28=EB=B3=80=EA=B8=B0=ED=98=B8=29?= Date: Fri, 14 May 2021 22:49:41 +0900 Subject: [PATCH] stage: allow duplicate values in _read_env (#6007) --- dvc/stage/__init__.py | 7 +++++-- tests/unit/stage/test_stage.py | 20 ++++++++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py index 4f89236f4..ae6a6e10b 100644 --- a/dvc/stage/__init__.py +++ b/dvc/stage/__init__.py @@ -296,8 +296,11 @@ class Stage(params.StageParams): for out in self.outs: current = self._read_env(out, checkpoint_func=checkpoint_func) - if set(env.keys()).intersection(set(current.keys())): - raise DvcException("Duplicated env variable") + if any( + env.get(key) != current.get(key) + for key in set(env.keys()).intersection(current.keys()) + ): + raise DvcException("Conflicting values for env variable") env.update(current) return env diff --git a/tests/unit/stage/test_stage.py b/tests/unit/stage/test_stage.py index 28135a5ec..0ad380ff6 100644 --- a/tests/unit/stage/test_stage.py +++ b/tests/unit/stage/test_stage.py @@ -122,3 +122,23 @@ def test_external_outs(tmp_path_factory, dvc): create_stage(Stage, dvc, "path.dvc", outs=["remote://myremote/foo"]) create_stage(Stage, dvc, "path.dvc", outs=[os.fspath(foo)], external=True) + + +def test_env(dvc): + from dvc.env import DVC_ROOT + from dvc.exceptions import DvcException + from dvc.stage import create_stage + + stage = create_stage(Stage, dvc, "path.dvc", outs=["foo", "bar"]) + + with mock.patch.object(stage, "_read_env", return_value={"foo": "foo"}): + env = stage.env() + assert env == {DVC_ROOT: dvc.root_dir, "foo": "foo"} + + def mock_read_env(out, **kwargs): + return {"foo": str(out)} + + with pytest.raises(DvcException) as exc: + with mock.patch.object(stage, "_read_env", mock_read_env): + env = stage.env() + assert exc.value == "Conflicting values for env variable"