diff --git a/tests/func/experiments/test_checkpoints.py b/tests/func/experiments/test_checkpoints.py index 35d6fa8ab..302ffbfbe 100644 --- a/tests/func/experiments/test_checkpoints.py +++ b/tests/func/experiments/test_checkpoints.py @@ -25,9 +25,7 @@ def test_new_checkpoint(tmp_dir, scm, dvc, checkpoint_stage, mocker, workspace, exp = first(results) new_mock.assert_called_once() - for rev in dvc.brancher([exp]): - if rev == "workspace": - continue + with dvc.switch(exp): fs = dvc.dvcfs with fs.open("foo") as fobj: assert fobj.read().strip() == str(checkpoint_stage.iterations) @@ -69,9 +67,7 @@ def test_resume_checkpoint( ) exp = first(results) - for rev in dvc.brancher([exp]): - if rev == "workspace": - continue + with dvc.switch(exp): fs = dvc.dvcfs with fs.open("foo") as fobj: assert fobj.read().strip() == str(2 * checkpoint_stage.iterations) @@ -94,9 +90,7 @@ def test_reset_checkpoint(tmp_dir, scm, dvc, checkpoint_stage, caplog, workspace ) exp = first(results) - for rev in dvc.brancher([exp]): - if rev == "workspace": - continue + with dvc.switch(exp): fs = dvc.dvcfs with fs.open("foo") as fobj: assert fobj.read().strip() == str(checkpoint_stage.iterations) @@ -132,18 +126,14 @@ def test_resume_branch(tmp_dir, scm, dvc, checkpoint_stage, workspace): ) checkpoint_b = first(results) - for rev in dvc.brancher([checkpoint_a]): - if rev == "workspace": - continue + with dvc.switch(checkpoint_a): fs = dvc.dvcfs with fs.open("foo") as fobj: assert fobj.read().strip() == str(2 * checkpoint_stage.iterations) with fs.open("metrics.yaml") as fobj: assert fobj.read().strip() == "foo: 2" - for rev in dvc.brancher([checkpoint_b]): - if rev == "workspace": - continue + with dvc.switch(checkpoint_b): fs = dvc.dvcfs with fs.open("foo") as fobj: assert fobj.read().strip() == str(2 * checkpoint_stage.iterations) diff --git a/tests/unit/fs/test_data.py b/tests/unit/fs/test_data.py index 20ab7a14f..180091b89 100644 --- a/tests/unit/fs/test_data.py +++ b/tests/unit/fs/test_data.py @@ -89,10 +89,7 @@ def test_open_in_history(tmp_dir, scm, dvc): dvc.scm.add(["foo.dvc", ".gitignore"]) dvc.scm.commit("foofoo") - for rev in dvc.brancher(revs=["HEAD~1"]): - if rev == "workspace": - continue - + with dvc.switch("HEAD~1"): fs = DataFileSystem(index=dvc.index.data["repo"]) with fs.open("foo", "r") as fobj: assert fobj.read() == "foo" diff --git a/tests/unit/fs/test_dvc.py b/tests/unit/fs/test_dvc.py index cf826e050..dbe68b805 100644 --- a/tests/unit/fs/test_dvc.py +++ b/tests/unit/fs/test_dvc.py @@ -59,10 +59,7 @@ def test_open_in_history(tmp_dir, scm, dvc): dvc.scm.add(["foo.dvc", ".gitignore"]) dvc.scm.commit("foofoo") - for rev in dvc.brancher(revs=["HEAD~1"]): - if rev == "workspace": - continue - + with dvc.switch("HEAD~1"): fs = DVCFileSystem(repo=dvc) with fs.open("foo", "r") as fobj: assert fobj.read() == "foo"