refactor(tests): standardize tests to follow pytest conventions

metaflow/plugins/pypi/pip.py

@@ -37,7 +37,7 @@ def __init__(self, error):
             self.package_spec = re.search(
                 "ERROR: No matching distribution found for (.*)", self.error
             )[1]
-            self.package_name = re.match("\w*", self.package_spec)[0]
+            self.package_name = re.match(r"\w*", self.package_spec)[0]
         except Exception:
             pass
 

test/cmd/diff/test_metaflow_diff.py

@@ -1,8 +1,6 @@
 import os
 import pytest
-import tempfile
 from subprocess import PIPE
-from unittest.mock import MagicMock, patch
 
 from metaflow.cmd.code import (
     extract_code_package,
@@ -14,131 +12,149 @@
 )
 
 
-class TestMetaflowDiff:
-
-    @patch("metaflow.cmd.code.Run")
-    def test_extract_code_package(self, mock_run):
-        mock_run.return_value.code.tarball.getmembers.return_value = []
-        mock_run.return_value.code.tarball.extractall = MagicMock()
-        runspec = "HelloFlow/3"
-
-        with patch(
-            "tempfile.TemporaryDirectory", return_value=tempfile.TemporaryDirectory()
-        ) as mock_tmp:
-            tmp = extract_code_package(runspec)
-
-        mock_run.assert_called_once_with(runspec, _namespace_check=False)
-        assert os.path.exists(tmp.name)
-
-    @pytest.mark.parametrize("use_tty", [True, False])
-    @patch("metaflow.cmd.code.run")
-    @patch("sys.stdout.isatty")
-    def test_perform_diff_output_false(self, mock_isatty, mock_run, use_tty):
-        mock_isatty.return_value = use_tty
-
-        mock_process = MagicMock()
-        mock_process.returncode = 1
-        mock_process.stdout = (
-            "--- a/file.txt\n"
-            "+++ b/file.txt\n"
-            "@@ -1 +1 @@\n"
-            "-source content\n"
-            "+target content\n"
-        )
-        mock_run.return_value = mock_process
-
-        with tempfile.TemporaryDirectory() as source_dir, tempfile.TemporaryDirectory() as target_dir:
-            source_file = os.path.join(source_dir, "file.txt")
-            target_file = os.path.join(target_dir, "file.txt")
-            with open(source_file, "w") as f:
-                f.write("source content")
-            with open(target_file, "w") as f:
-                f.write("target content")
-
-            perform_diff(source_dir, target_dir, output=False)
-
-            # if output=False, run should be called twice:
-            # 1. git diff
-            # 2. less -R
-            assert mock_run.call_count == 2
-
-            mock_run.assert_any_call(
-                [
-                    "git",
-                    "diff",
-                    "--no-index",
-                    "--exit-code",
-                    "--color" if use_tty else "--no-color",
-                    "./file.txt",
-                    source_file,
-                ],
-                text=True,
-                stdout=PIPE,
-                cwd=target_dir,
-            )
-
-            mock_run.assert_any_call(
-                ["less", "-R"], input=mock_process.stdout, text=True
-            )
-
-    @patch("metaflow.cmd.code.run")
-    def test_perform_diff_output_true(self, mock_run):
-        with tempfile.TemporaryDirectory() as source_dir, tempfile.TemporaryDirectory() as target_dir:
-            source_file = os.path.join(source_dir, "file.txt")
-            target_file = os.path.join(target_dir, "file.txt")
-            with open(source_file, "w") as f:
-                f.write("source content")
-            with open(target_file, "w") as f:
-                f.write("target content")
-
-            perform_diff(source_dir, target_dir, output=True)
-
-            assert mock_run.call_count == 1
-
-            mock_run.assert_called_once_with(
-                [
-                    "git",
-                    "diff",
-                    "--no-index",
-                    "--exit-code",
-                    "--no-color",
-                    "./file.txt",
-                    source_file,
-                ],
-                text=True,
-                stdout=PIPE,
-                cwd=target_dir,
-            )
-
-    @patch("shutil.rmtree")
-    @patch("metaflow.cmd.code.extract_code_package")
-    @patch("metaflow.cmd.code.op_diff")
-    def test_run_op(self, mock_op_diff, mock_extract_code_package, mock_rmtree):
-        mock_tmp = tempfile.TemporaryDirectory()
-        mock_extract_code_package.return_value = mock_tmp
-        runspec = "HelloFlow/3"
-
-        run_op(runspec, mock_op_diff)
-
-        mock_extract_code_package.assert_called_once_with(runspec)
-        mock_op_diff.assert_called_once_with(mock_tmp.name)
-
-        mock_rmtree.assert_any_call(mock_tmp.name)
-
-    @patch("metaflow.cmd.code.perform_diff")
-    def test_op_patch(self, mock_perform_diff):
-        mock_perform_diff.return_value = ["diff --git a/file.txt b/file.txt\n"]
-
-        with tempfile.TemporaryDirectory() as tmpdir:
-            patch_file = os.path.join(tmpdir, "patch.patch")
-
-            op_patch(tmpdir, patch_file)
-
-            mock_perform_diff.assert_called_once_with(tmpdir, output=True)
-            with open(patch_file, "r") as f:
-                content = f.read()
-            assert "diff --git a/file.txt b/file.txt\n" in content
-
-
-if __name__ == "__main__":
-    pytest.main([__file__])
+def test_extract_code_package_creates_temp_dir(mocker):
+    """Test that extract_code_package safely unpacks the tarball into a temporary directory."""
+
+    mock_run = mocker.patch("metaflow.cmd.code.Run")
+    mock_run.return_value.code.tarball.getmembers.return_value = []
+    mock_run.return_value.code.tarball.extractall = mocker.MagicMock()
+
+    mock_tmp = mocker.patch("tempfile.TemporaryDirectory")
+    mock_tmp.return_value.name = "/fake/tmp/dir"
+
+    runspec = "HelloFlow/3"
+
+    # Act
+    tmp = extract_code_package(runspec)
+
+    # Assert
+    mock_run.assert_called_once_with(runspec, _namespace_check=False)
+    assert tmp.name == "/fake/tmp/dir"
+
+
+@pytest.mark.parametrize(
+    "use_tty", [True, False], ids=["with_tty_colors", "without_tty_colors"]
+)
+def test_perform_diff_triggers_git_and_less_when_output_false(
+    mocker, tmp_path, use_tty
+):
+    """Test that perform_diff executes git diff and pipes stdout to less when output is False."""
+    mock_isatty = mocker.patch("sys.stdout.isatty", return_value=use_tty)
+    mock_run = mocker.patch("metaflow.cmd.code.run")
+
+    # Setup: Create a fake git diff process output
+    mock_process = mocker.MagicMock()
+    mock_process.returncode = 1
+    mock_process.stdout = (
+        "--- a/file.txt\n"
+        "+++ b/file.txt\n"
+        "@@ -1 +1 @@\n"
+        "-source content\n"
+        "+target content\n"
+    )
+    mock_run.return_value = mock_process
+
+    # Setup: Create physical files using pytest's tmp_path fixture
+    source_dir = tmp_path / "source"
+    target_dir = tmp_path / "target"
+    source_dir.mkdir()
+    target_dir.mkdir()
+
+    source_file = source_dir / "file.txt"
+    target_file = target_dir / "file.txt"
+    source_file.write_text("source content")
+    target_file.write_text("target content")
+
+    # Act
+    perform_diff(str(source_dir), str(target_dir), output=False)
+
+    # Assert: git diff and less should both be called sequentially
+    assert mock_run.call_count == 2
+    mock_run.assert_any_call(
+        [
+            "git",
+            "diff",
+            "--no-index",
+            "--exit-code",
+            "--color" if use_tty else "--no-color",
+            "./file.txt",
+            str(source_file),
+        ],
+        text=True,
+        stdout=PIPE,
+        cwd=str(target_dir),
+    )
+    mock_run.assert_any_call(["less", "-R"], input=mock_process.stdout, text=True)
+
+
+def test_perform_diff_returns_early_when_output_true(mocker, tmp_path):
+    """Test that perform_diff only runs git diff and skips less when output is True."""
+    mock_run = mocker.patch("metaflow.cmd.code.run")
+
+    # Setup: Create physical files
+    source_dir = tmp_path / "source"
+    target_dir = tmp_path / "target"
+    source_dir.mkdir()
+    target_dir.mkdir()
+
+    source_file = source_dir / "file.txt"
+    target_file = target_dir / "file.txt"
+    source_file.write_text("source content")
+    target_file.write_text("target content")
+
+    # Act
+    perform_diff(str(source_dir), str(target_dir), output=True)
+
+    # Assert: Only git diff is called
+    assert mock_run.call_count == 1
+    mock_run.assert_called_once_with(
+        [
+            "git",
+            "diff",
+            "--no-index",
+            "--exit-code",
+            "--no-color",
+            "./file.txt",
+            str(source_file),
+        ],
+        text=True,
+        stdout=PIPE,
+        cwd=str(target_dir),
+    )
+
+
+def test_run_op_cleans_up_temporary_directory_after_execution(mocker):
+    """Test that run_op delegates to op_diff and correctly tears down the temp directory."""
+    mock_rmtree = mocker.patch("shutil.rmtree")
+    mock_extract = mocker.patch("metaflow.cmd.code.extract_code_package")
+    mock_op_diff = mocker.MagicMock()
+
+    # Setup: Mock the temporary directory object returned by extract_code_package
+    mock_tmp = mocker.MagicMock()
+    mock_tmp.name = "/fake/tmp/dir"
+    mock_extract.return_value = mock_tmp
+
+    runspec = "HelloFlow/3"
+
+    # Act
+    run_op(runspec, mock_op_diff)
+
+    # Assert
+    mock_extract.assert_called_once_with(runspec)
+    mock_op_diff.assert_called_once_with("/fake/tmp/dir")
+    mock_rmtree.assert_any_call("/fake/tmp/dir")
+
+
+def test_op_patch_writes_diff_output_to_file(mocker, tmp_path):
+    """Test that op_patch successfully writes the generated diff to the specified patch file."""
+    mock_perform_diff = mocker.patch("metaflow.cmd.code.perform_diff")
+    mock_perform_diff.return_value = ["diff --git a/file.txt b/file.txt\n"]
+
+    patch_file = tmp_path / "patch.patch"
+
+    # Act
+    op_patch(str(tmp_path), str(patch_file))
+
+    # Assert
+    mock_perform_diff.assert_called_once_with(str(tmp_path), output=True)
+    assert patch_file.read_text() == "diff --git a/file.txt b/file.txt\n"

test/data/s3/conftest.py

@@ -30,9 +30,8 @@ def reset_current_env():
     Fixture to ensure the metaflow current environment is clean between tests.
 
     This prevents test pollution when tests manipulate the global current state.
-    The fixture runs automatically for all tests in this directory.
     """
-    # Setup: Save all private attributes that might be set by current._set_env
+    # Setup: Capture internal state
     saved_state = {
         attr: getattr(current, attr, None)
         for attr in dir(current)
@@ -41,27 +40,21 @@ def reset_current_env():
 
     yield
 
-    # Teardown: Clear all current environment attributes
-    # First, remove any new attributes that were added
-    for attr in dir(current):
-        if (
-            attr.startswith("_")
-            and not attr.startswith("__")
-            and attr not in saved_state
-        ):
-            try:
-                delattr(current, attr)
-            except AttributeError:
-                pass  # Some attributes may be read-only
+    # Teardown: Restore original attributes and clean up new ones
+    current_attrs = [
+        attr
+        for attr in dir(current)
+        if attr.startswith("_") and not attr.startswith("__")
+    ]
 
-    # Then restore original values
-    for attr, value in saved_state.items():
+    for attr in current_attrs:
         try:
-            if value is None:
-                # Remove attribute if it didn't exist before
-                if hasattr(current, attr):
-                    delattr(current, attr)
+            if attr not in saved_state:
+                # Remove attributes created during the test
+                delattr(current, attr)
             else:
-                setattr(current, attr, value)
-        except AttributeError:
-            pass  # Some attributes may be read-only
+                # Restore original values
+                setattr(current, attr, saved_state[attr])
+        except (AttributeError, TypeError):
+            # Some internal attributes in 'current' may be read-only or immutable
+            pass