fix tests

tests/test_metaflow_diff.py

@@ -1,6 +1,6 @@
 import os
+import pytest
 import tempfile
-import unittest
 from subprocess import PIPE
 from unittest.mock import MagicMock, patch
 
@@ -15,7 +15,7 @@
 )
 
 
-class TestMetaflowDiff(unittest.TestCase):
+class TestMetaflowDiff:
 
     @patch("metaflow_diff.metaflow_diff.Run")
     @patch("metaflow_diff.metaflow_diff.namespace")
@@ -31,10 +31,61 @@ def test_extract_code_package(self, mock_namespace, mock_run):
 
         mock_namespace.assert_called_once_with(None)
         mock_run.assert_called_once_with(runspec)
-        self.assertTrue(os.path.exists(tmp.name))
+        assert os.path.exists(tmp.name)
 
+    @pytest.mark.parametrize("use_tty", [True, False])
     @patch("metaflow_diff.metaflow_diff.run")
-    def test_perform_diff(self, mock_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_diff.metaflow_diff.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")
@@ -43,10 +94,23 @@ def test_perform_diff(self, mock_run):
             with open(target_file, "w") as f:
                 f.write("target content")
 
-            perform_diff(source_dir, target_dir)
+            perform_diff(source_dir, target_dir, output=True)
+
+            assert mock_run.call_count == 1
 
             mock_run.assert_called_once_with(
-                ["git", "diff", "--no-index", source_file, target_file],
+                [
+                    "git",
+                    "diff",
+                    "--no-index",
+                    "--exit-code",
+                    "--no-color",
+                    "./file.txt",
+                    source_file,
+                ],
+                text=True,
+                stdout=PIPE,
+                cwd=target_dir,
             )
 
     @patch("shutil.rmtree")
@@ -61,7 +125,8 @@ def test_run_op(self, mock_op_diff, mock_extract_code_package, mock_rmtree):
 
         mock_extract_code_package.assert_called_once_with(runspec, EXCLUSIONS)
         mock_op_diff.assert_called_once_with(mock_tmp.name)
-        mock_rmtree.assert_called_once_with(mock_tmp.name)
+
+        mock_rmtree.assert_any_call(mock_tmp.name)
 
     @patch("metaflow_diff.metaflow_diff.perform_diff")
     def test_op_patch(self, mock_perform_diff):
@@ -72,11 +137,11 @@ def test_op_patch(self, mock_perform_diff):
 
             op_patch(tmpdir, patch_file)
 
-            mock_perform_diff.assert_called_once_with(tmpdir, os.getcwd(), output=True)
+            mock_perform_diff.assert_called_once_with(tmpdir, output=True)
             with open(patch_file, "r") as f:
                 content = f.read()
-            self.assertIn("diff --git a/file.txt b/file.txt\n", content)
+            assert "diff --git a/file.txt b/file.txt\n" in content
 
 
 if __name__ == "__main__":
-    unittest.main()
+    pytest.main([__file__])