Refactored and improved the tests

Author: tim-u

google/cloud/dataproc_magics/magics.py

@@ -29,26 +29,28 @@ def __init__(
     ):
         super().__init__(shell, **kwargs)
 
-    def _parse_command(self, args):
-        if not args or args[0] != "install":
-            print("Usage: %dpip install <package1> <package2> ...")
-            return
-
-        # filter out 'install' and the flags (not currently supported)
-        packages = [pkg for pkg in args[1:] if not pkg.startswith("-")]
-        return packages
-
     @line_magic
     def dpip(self, line):
         """
         Custom magic to install pip packages as Spark Connect artifacts.
         Usage: %dpip install pandas numpy
         """
         try:
-            packages = self._parse_command(shlex.split(line))
+            args = shlex.split(line)
+
+            if not args or args[0].lower() != "install":
+                print("Usage: %dpip install <package1> <package2> ...")
+                return
+
+            packages = args[1:]  # remove `install`
 
             if not packages:
-                print("No packages specified.")
+                print("Error: No packages specified.")
+                return
+
+            # 4. Check for unsupported flags
+            if any(pkg.startswith("-") for pkg in packages):
+                print("Error: Flags are not currently supported.")
                 return
 
             sessions = [
@@ -59,15 +61,16 @@ def dpip(self, line):
 
             if not sessions:
                 print(
-                    "No active Spark Sessions found. Please create one first."
+                    "No active Dataproc Spark Sessions found. Please create one first."
                 )
                 return
 
+            print("Active sessions found: %s", self.shell.user_ns.keys())
             print("Installing packages: %s", packages)
             for session in sessions:
                 for package in packages:
                     session.addArtifacts(package, pypi=True)
 
-            print("Packages successfully added as artifacts.")
+            print("Successfully installed packages in Dataproc session(s).")
         except Exception as e:
-            print(f"Failed to add artifacts: {e}")
+            print(f"Failed to install packages: {e}")

tests/integration/dataproc_magics/test_magics.py

@@ -68,7 +68,8 @@ def test_subnet():
 
 @pytest.fixture
 def test_subnetwork_uri(test_subnet):
-    # Make DATAPROC_SPARK_CONNECT_SUBNET the full URI to align with how user would specify it in the project
+    # Make DATAPROC_SPARK_CONNECT_SUBNET the full URI
+    # to align with how user would specify it in the project
     return test_subnet
 
 
@@ -106,11 +107,9 @@ def connect_session(test_project, test_region, os_environment):
         pass
 
 
-# Tests for magics.py
 @pytest.fixture
 def ipython_shell(connect_session):
     """Provides an IPython shell with a DataprocSparkSession in user_ns."""
-    pytest.importorskip("IPython", reason="IPython not available")
     try:
         from IPython.terminal.interactiveshell import TerminalInteractiveShell
         from google.cloud import dataproc_magics
@@ -128,71 +127,79 @@ def ipython_shell(connect_session):
         TerminalInteractiveShell.clear_instance()
 
 
+# Tests for magics.py
 def test_dpip_magic_loads(ipython_shell):
     """Test that %dpip magic is registered."""
     assert "dpip" in ipython_shell.magics_manager.magics["line"]
 
 
-@mock.patch.object(DataprocSparkSession, "addArtifacts")
-def test_dpip_install_single_package(mock_add_artifacts, ipython_shell, capsys):
+def test_dpip_install_success(connect_session, ipython_shell, capsys):
     """Test installing a single package with %dpip."""
-    ipython_shell.run_line_magic("dpip", "install pandas")
-    mock_add_artifacts.assert_called_once_with("pandas", pypi=True)
+    ipython_shell.run_line_magic("dpip", "install roman")
     captured = capsys.readouterr()
-    assert "Installing packages: " in captured.out
-    assert "Packages successfully added as artifacts." in captured.out
+    assert "Active sessions found:" in captured.out
+    assert "Installing packages:" in captured.out
+    assert (
+        "Successfully installed packages in Dataproc session(s)."
+        in captured.out
+    )
 
+    from pyspark.sql.connect.functions import udf
+    from pyspark.sql.types import StringType
 
-@mock.patch.object(DataprocSparkSession, "addArtifacts")
-def test_dpip_install_multiple_packages_with_flags(
-    mock_add_artifacts, ipython_shell, capsys
-):
-    """Test installing multiple packages with flags like -U."""
-    ipython_shell.run_line_magic("dpip", "install -U numpy scikit-learn")
-    calls = [
-        mock.call("numpy", pypi=True),
-        mock.call("scikit-learn", pypi=True),
-    ]
-    mock_add_artifacts.assert_has_calls(calls, any_order=True)
-    assert mock_add_artifacts.call_count == 2
-    captured = capsys.readouterr()
-    assert "Installing packages: " in captured.out
-    assert "Packages successfully added as artifacts." in captured.out
+    df = connect_session.createDataFrame(
+        [(1,), (4,), (16,), (51,), (1666,)], ["number"]
+    )
+
+    def to_roman(number):
+        import roman
+
+        return roman.toRoman(number) if number else None
+
+    df_result = df.withColumn(
+        "roman", udf(to_roman, StringType())("number")
+    ).collect()
+
+    assert df_result[0]["roman"] == "I"
+    assert df_result[1]["roman"] == "IV"
+    assert df_result[2]["roman"] == "XVI"
+    assert df_result[3]["roman"] == "LI"
+    assert df_result[4]["roman"] == "MDCLXVI"
+
+    connect_session.stop()
 
 
 def test_dpip_no_install_command(ipython_shell, capsys):
     """Test usage message when 'install' is missing."""
     ipython_shell.run_line_magic("dpip", "pandas")
     captured = capsys.readouterr()
     assert "Usage: %dpip install <package1> <package2> ..." in captured.out
-    assert "No packages specified." in captured.out
 
 
 def test_dpip_no_packages(ipython_shell, capsys):
     """Test message when no packages are specified."""
     ipython_shell.run_line_magic("dpip", "install")
     captured = capsys.readouterr()
-    assert "No packages specified." in captured.out
+    assert "Error: No packages specified." in captured.out
+
+
+def test_dpip_with_flags(ipython_shell, capsys):
+    """Test installing multiple packages with flags like -U."""
+    ipython_shell.run_line_magic("dpip", "install -U numpy scikit-learn")
+    captured = capsys.readouterr()
+    assert "Error: Flags are not currently supported." in captured.out
 
 
-@mock.patch.object(DataprocSparkSession, "addArtifacts")
-def test_dpip_no_session(mock_add_artifacts, ipython_shell, capsys):
+def test_dpip_no_session(ipython_shell, capsys):
     """Test message when no Spark session is active."""
     ipython_shell.user_ns = {}  # Remove spark session from namespace
     ipython_shell.run_line_magic("dpip", "install pandas")
     captured = capsys.readouterr()
-    assert "No active Spark Sessions found." in captured.out
-    mock_add_artifacts.assert_not_called()
+    assert "No active Dataproc Spark Sessions found." in captured.out
 
 
-@mock.patch.object(
-    DataprocSparkSession,
-    "addArtifacts",
-    side_effect=Exception("Install failed"),
-)
-def test_dpip_install_failure(mock_add_artifacts, ipython_shell, capsys):
+def test_dpip_install_failure(ipython_shell, capsys):
     """Test error message on installation failure."""
-    ipython_shell.run_line_magic("dpip", "install bad-package")
-    mock_add_artifacts.assert_called_once_with("bad-package", pypi=True)
+    ipython_shell.run_line_magic("dpip", "install dp-non-existent-package")
     captured = capsys.readouterr()
-    assert "Failed to add artifacts: Install failed" in captured.out
+    assert "No matching distribution found" in captured.out

tests/unit/dataproc_magics/test_magics.py

@@ -31,39 +31,39 @@ def setUp(self):
         self.shell.config = Config()
         self.magics = DataprocMagics(shell=self.shell)
 
-    def test_parse_command_valid(self):
-        packages = self.magics._parse_command(["install", "pandas", "numpy"])
-        self.assertEqual(packages, ["pandas", "numpy"])
+    def test_dpip_with_flags(self):
+        f = io.StringIO()
+        with redirect_stdout(f):
+            self.magics.dpip("install --upgrade numpy")
+        self.assertIn("Error: Flags are not currently supported.", f.getvalue())
 
-    def test_parse_command_with_flags(self):
-        packages = self.magics._parse_command(
-            ["install", "-U", "pandas", "--upgrade", "numpy"]
+    def test_dpip_no_install(self):
+        f = io.StringIO()
+        with redirect_stdout(f):
+            self.magics.dpip("pandas numpy")
+        self.assertIn(
+            "Usage: %dpip install <package1> <package2> ...", f.getvalue()
         )
-        self.assertEqual(packages, ["pandas", "numpy"])
-
-    def test_parse_command_no_install(self):
-        packages = self.magics._parse_command(["other", "pandas"])
-        self.assertIsNone(packages)
 
     def test_dpip_invalid_command(self):
         f = io.StringIO()
         with redirect_stdout(f):
             self.magics.dpip("foo bar")
-        output = f.getvalue()
-        self.assertIn("Usage: %dpip install", output)
-        self.assertIn("No packages specified", output)
+        self.assertIn(
+            "Usage: %dpip install <package1> <package2> ...", f.getvalue()
+        )
 
     def test_dpip_no_session(self):
         f = io.StringIO()
         with redirect_stdout(f):
             self.magics.dpip("install pandas")
-        self.assertIn("No active Spark Sessions found", f.getvalue())
+        self.assertIn("No active Dataproc Spark Sessions found", f.getvalue())
 
     def test_dpip_no_packages_specified(self):
         f = io.StringIO()
         with redirect_stdout(f):
             self.magics.dpip("install")
-        self.assertIn("No packages specified", f.getvalue())
+        self.assertIn("Error: No packages specified", f.getvalue())
 
     def test_dpip_install_packages_single_session(self):
         mock_session = mock.Mock(spec=DataprocSparkSession)
@@ -80,7 +80,10 @@ def test_dpip_install_packages_single_session(self):
             ]
         )
         self.assertEqual(mock_session.addArtifacts.call_count, 2)
-        self.assertIn("Packages successfully added as artifacts.", f.getvalue())
+        self.assertIn(
+            "Successfully installed packages in Dataproc session(s).",
+            f.getvalue(),
+        )
 
     def test_dpip_install_packages_multiple_sessions(self):
         mock_session1 = mock.Mock(spec=DataprocSparkSession)
@@ -95,7 +98,10 @@ def test_dpip_install_packages_multiple_sessions(self):
 
         mock_session1.addArtifacts.assert_called_once_with("pandas", pypi=True)
         mock_session2.addArtifacts.assert_called_once_with("pandas", pypi=True)
-        self.assertIn("Packages successfully added as artifacts.", f.getvalue())
+        self.assertIn(
+            "Successfully installed packages in Dataproc session(s).",
+            f.getvalue(),
+        )
 
     def test_dpip_add_artifacts_fails(self):
         mock_session = mock.Mock(spec=DataprocSparkSession)
@@ -107,18 +113,7 @@ def test_dpip_add_artifacts_fails(self):
             self.magics.dpip("install pandas")
 
         mock_session.addArtifacts.assert_called_once_with("pandas", pypi=True)
-        self.assertIn("Failed to add artifacts: Failed", f.getvalue())
-
-    def test_dpip_with_flags(self):
-        mock_session = mock.Mock(spec=DataprocSparkSession)
-        self.shell.user_ns["spark"] = mock_session
-
-        f = io.StringIO()
-        with redirect_stdout(f):
-            self.magics.dpip("install -U pandas")
-
-        mock_session.addArtifacts.assert_called_once_with("pandas", pypi=True)
-        self.assertIn("Packages successfully added as artifacts.", f.getvalue())
+        self.assertIn("Failed to install packages: Failed", f.getvalue())
 
 
 if __name__ == "__main__":