Unit test updates

Author: markdlv-sf

tests/test_deploy.py

@@ -5,6 +5,7 @@
     mock_open,
     patch,
 )
+import zipfile
 
 import pytest
 import requests
@@ -28,9 +29,10 @@
         get_data_transform_config,
         get_deployments,
         run_data_transform,
+        upload_zip,
         verify_data_transform_config,
         wait_for_deployment,
-        zip_and_upload_directory,
+        zip,
     )
 
 
@@ -134,22 +136,100 @@ def test_create_deployment_conflict(self, mock_make_api_call):
             create_deployment(access_token, metadata)
 
 
-class TestZipAndUploadDirectory:
-    @patch("datacustomcode.deploy.shutil.make_archive")
+class TestZip:
+    @patch("datacustomcode.deploy.has_nonempty_requirements_file")
+    @patch("datacustomcode.deploy.prepare_dependency_archive")
+    @patch("zipfile.ZipFile")
+    @patch("os.walk")
+    def test_zip_with_requirements(
+        self, mock_walk, mock_zipfile, mock_prepare, mock_has_requirements
+    ):
+        """Test zipping a directory with requirements.txt."""
+        mock_has_requirements.return_value = True
+        mock_zipfile_instance = MagicMock()
+        mock_zipfile.return_value.__enter__.return_value = mock_zipfile_instance
+        mock_zipfile_instance.write = MagicMock()
+
+        # Mock os.walk to return some test files
+        mock_walk.return_value = [
+            ("/test/dir", ["subdir"], ["file1.py", "file2.py"]),
+            ("/test/dir/subdir", [], ["file3.py"]),
+        ]
+
+        zip("/test/dir")
+
+        mock_has_requirements.assert_called_once_with("/test/dir")
+        mock_prepare.assert_called_once_with("/test/dir")
+        mock_zipfile.assert_called_once_with(
+            "deployment.zip", "w", zipfile.ZIP_DEFLATED
+        )
+        assert mock_zipfile_instance.write.call_count == 3  # One call per file
+
+    @patch("datacustomcode.deploy.has_nonempty_requirements_file")
+    @patch("datacustomcode.deploy.prepare_dependency_archive")
+    @patch("zipfile.ZipFile")
+    @patch("os.walk")
+    def test_zip_without_requirements(
+        self, mock_walk, mock_zipfile, mock_prepare, mock_has_requirements
+    ):
+        """Test zipping a directory without requirements.txt."""
+        mock_has_requirements.return_value = False
+        mock_zipfile_instance = MagicMock()
+        mock_zipfile.return_value.__enter__.return_value = mock_zipfile_instance
+        mock_zipfile_instance.write = MagicMock()
+
+        # Mock os.walk to return some test files
+        mock_walk.return_value = [
+            ("/test/dir", ["subdir"], ["file1.py", "file2.py"]),
+            ("/test/dir/subdir", [], ["file3.py"]),
+        ]
+
+        zip("/test/dir")
+
+        mock_has_requirements.assert_called_once_with("/test/dir")
+        mock_prepare.assert_not_called()
+        mock_zipfile.assert_called_once_with(
+            "deployment.zip", "w", zipfile.ZIP_DEFLATED
+        )
+        assert mock_zipfile_instance.write.call_count == 3  # One call per file
+
+
+class TestUploadZip:
     @patch("datacustomcode.deploy.requests.put")
     @patch("builtins.open", new_callable=mock_open, read_data=b"test data")
-    def test_zip_and_upload_directory(self, mock_file, mock_put, mock_make_archive):
-        """Test zipping and uploading a directory."""
+    def test_upload_zip_success(self, mock_file, mock_put):
+        """Test successful zip upload."""
         mock_response = MagicMock()
         mock_put.return_value = mock_response
 
-        zip_and_upload_directory("/test/dir", "https://upload.example.com")
+        upload_zip("https://upload.example.com")
 
-        mock_make_archive.assert_called_once_with("deployment", "zip", "/test/dir")
         mock_file.assert_called_once_with("deployment.zip", "rb")
-        mock_put.assert_called_once()
+        mock_put.assert_called_once_with(
+            "https://upload.example.com",
+            data=mock_file.return_value,
+            headers={"Content-Type": "application/zip"},
+        )
         mock_response.raise_for_status.assert_called_once()
 
+    @patch("datacustomcode.deploy.requests.put")
+    @patch("builtins.open", new_callable=mock_open, read_data=b"test data")
+    def test_upload_zip_http_error(self, mock_file, mock_put):
+        """Test zip upload with HTTP error."""
+        mock_response = MagicMock()
+        mock_response.raise_for_status.side_effect = requests.HTTPError("Upload failed")
+        mock_put.return_value = mock_response
+
+        with pytest.raises(requests.HTTPError, match="Upload failed"):
+            upload_zip("https://upload.example.com")
+
+        mock_file.assert_called_once_with("deployment.zip", "rb")
+        mock_put.assert_called_once_with(
+            "https://upload.example.com",
+            data=mock_file.return_value,
+            headers={"Content-Type": "application/zip"},
+        )
+
 
 class TestGetDeployments:
     @patch("datacustomcode.deploy._make_api_call")
@@ -328,14 +408,16 @@ class TestDeployFull:
     @patch("datacustomcode.deploy.prepare_dependency_archive")
     @patch("datacustomcode.deploy.verify_data_transform_config")
     @patch("datacustomcode.deploy.create_deployment")
-    @patch("datacustomcode.deploy.zip_and_upload_directory")
+    @patch("datacustomcode.deploy.zip")
+    @patch("datacustomcode.deploy.upload_zip")
     @patch("datacustomcode.deploy.wait_for_deployment")
     @patch("datacustomcode.deploy.create_data_transform")
     def test_deploy_full(
         self,
         mock_create_transform,
         mock_wait,
-        mock_zip_upload,
+        mock_upload_zip,
+        mock_zip,
         mock_create_deployment,
         mock_verify_config,
         mock_prepare,
@@ -371,9 +453,8 @@ def test_deploy_full(
         mock_prepare.assert_called_once_with("/test/dir")
         mock_verify_config.assert_called_once_with("/test/dir")
         mock_create_deployment.assert_called_once_with(access_token, metadata)
-        mock_zip_upload.assert_called_once_with(
-            "/test/dir", "https://upload.example.com"
-        )
+        mock_zip.assert_called_once_with("/test/dir")
+        mock_upload_zip.assert_called_once_with("https://upload.example.com")
         mock_wait.assert_called_once_with(access_token, metadata, callback)
         mock_create_transform.assert_called_once_with(
             "/test/dir", access_token, metadata

tests/test_scan.py

@@ -11,6 +11,8 @@
     DataAccessLayerCalls,
     dc_config_json_from_file,
     scan_file,
+    scan_file_for_imports,
+    write_requirements_file,
 )
 
 
@@ -422,3 +424,130 @@ def process_customer_data():
         assert config["permissions"]["write"]["dlo"] == ["customer_data_enriched"]
     finally:
         os.unlink(temp_path)
+
+
+class TestRequirementsFile:
+    def test_scan_file_for_imports(self):
+        """Test scanning a file for external package imports."""
+        content = textwrap.dedent(
+            """
+            import pandas as pd
+            import numpy as np
+            from sklearn.linear_model import LinearRegression
+            import os  # Standard library
+            import sys  # Standard library
+            from datacustomcode.client import Client  # Internal package
+            """
+        )
+        temp_path = create_test_script(content)
+        try:
+            imports = scan_file_for_imports(temp_path)
+            assert "pandas" in imports
+            assert "numpy" in imports
+            assert "sklearn" in imports
+            assert "os" not in imports  # Standard library
+            assert "sys" not in imports  # Standard library
+            assert "datacustomcode" not in imports  # Internal package
+        finally:
+            os.unlink(temp_path)
+
+    def test_write_requirements_file_new(self):
+        """Test writing a new requirements.txt file."""
+        content = textwrap.dedent(
+            """
+            import pandas as pd
+            import numpy as np
+            """
+        )
+        temp_path = create_test_script(content)
+        try:
+            requirements_path = write_requirements_file(temp_path)
+            assert os.path.exists(requirements_path)
+
+            with open(requirements_path, "r") as f:
+                requirements = {line.strip() for line in f}
+
+            assert "pandas" in requirements
+            assert "numpy" in requirements
+        finally:
+            os.unlink(temp_path)
+            if os.path.exists(requirements_path):
+                os.unlink(requirements_path)
+
+    def test_write_requirements_file_merge(self):
+        """Test merging with existing requirements.txt file."""
+        # First create an existing requirements.txt
+        temp_dir = tempfile.mkdtemp()
+        existing_requirements = os.path.join(temp_dir, "requirements.txt")
+        with open(existing_requirements, "w") as f:
+            f.write("pandas\nnumpy\n")
+
+        # Create a new Python file with additional imports
+        content = textwrap.dedent(
+            """
+            import pandas as pd
+            import numpy as np
+            import scipy
+            import matplotlib
+            """
+        )
+        temp_path = create_test_script(content)
+        try:
+            requirements_path = write_requirements_file(temp_path)
+            assert os.path.exists(requirements_path)
+
+            with open(requirements_path, "r") as f:
+                requirements = {line.strip() for line in f}
+
+            # Check that both existing and new requirements are present
+            assert "pandas" in requirements
+            assert "numpy" in requirements
+            assert "scipy" in requirements
+            assert "matplotlib" in requirements
+        finally:
+            os.unlink(temp_path)
+            if os.path.exists(requirements_path):
+                os.unlink(requirements_path)
+            if os.path.exists(existing_requirements):
+                os.unlink(existing_requirements)
+            os.rmdir(temp_dir)
+
+    def test_standard_library_exclusion(self):
+        """Test that standard library imports are excluded."""
+        content = textwrap.dedent(
+            """
+            import os
+            import sys
+            import json
+            import datetime
+            import pandas as pd
+            """
+        )
+        temp_path = create_test_script(content)
+        try:
+            imports = scan_file_for_imports(temp_path)
+            assert "pandas" in imports
+            assert "os" not in imports
+            assert "sys" not in imports
+            assert "json" not in imports
+            assert "datetime" not in imports
+        finally:
+            os.unlink(temp_path)
+
+    def test_excluded_packages(self):
+        """Test that excluded packages are not included in requirements."""
+        content = textwrap.dedent(
+            """
+            import datacustomcode
+            import pyspark
+            import pandas as pd
+            """
+        )
+        temp_path = create_test_script(content)
+        try:
+            imports = scan_file_for_imports(temp_path)
+            assert "pandas" in imports
+            assert "datacustomcode" not in imports
+            assert "pyspark" not in imports
+        finally:
+            os.unlink(temp_path)