feat(sql_connectors): add unit tests for sql connectors

tests/unit_tests/dataframe/test_dataframe.py

@@ -1,11 +1,13 @@
-from unittest.mock import Mock, patch
+from unittest.mock import MagicMock, Mock, mock_open, patch
 
 import pandas as pd
 import pytest
+from numpy import False_
 
 import pandasai
 from pandasai.agent import Agent
 from pandasai.dataframe.base import DataFrame
+from pandasai.exceptions import PandaAIApiKeyError
 
 
 class TestDataFrame:
@@ -26,7 +28,12 @@ def sample_data(self):
 
     @pytest.fixture
     def sample_df(self, sample_data):
-        return DataFrame(sample_data)
+        return DataFrame(
+            sample_data,
+            path="acme-corp/employees",
+            name="employees",
+            description="Employee data",
+        )
 
     def test_dataframe_initialization(self, sample_data, sample_df):
         assert isinstance(sample_df, DataFrame)
@@ -88,3 +95,109 @@ def test_column_hash(self, sample_df):
         assert hasattr(sample_df, "column_hash")
         assert isinstance(sample_df.column_hash, str)
         assert len(sample_df.column_hash) == 32  # MD5 hash length
+
+    @patch("pandasai.dataframe.base.get_pandaai_session")
+    @patch("pandasai.dataframe.base.os.path.exists")
+    @patch("pandasai.dataframe.base.open", new_callable=mock_open)
+    @patch("pandasai.dataframe.base.os.environ")
+    @patch("pandasai.dataframe.base.find_project_root")
+    def test_push_successful(
+        self,
+        mock_find_project_root,
+        mock_environ,
+        mock_open,
+        mock_path_exists,
+        mock_get_session,
+        sample_df,
+    ):
+        # Mock environment variable
+        mock_environ.get.return_value = "fake_api_key"
+
+        # Mock project root
+        mock_find_project_root.return_value = "/fake/project/root"
+
+        # Mock file paths
+        mock_path_exists.side_effect = (
+            lambda x: "data.parquet" in x
+        )  # Only data.parquet exists
+
+        # Mock session and POST request
+        mock_session = MagicMock()
+        mock_get_session.return_value = mock_session
+
+        sample_df.push()
+
+        # Verify POST request
+        mock_session.post.assert_called_once_with(
+            "/datasets/push",
+            files=[
+                (
+                    "files",
+                    ("schema.yaml", mock_open.return_value, "application/x-yaml"),
+                ),
+                (
+                    "files",
+                    (
+                        "data.parquet",
+                        mock_open.return_value,
+                        "application/octet-stream",
+                    ),
+                ),
+            ],
+            params={
+                "path": sample_df.path,
+                "description": sample_df.description,
+                "name": sample_df.name,
+            },
+            headers={
+                "accept": "application/json",
+                "x-authorization": "Bearer fake_api_key",
+            },
+        )
+
+    def test_push_raises_error_if_path_is_none(self, sample_df):
+        # Set up the object with no path
+        sample_df.path = None
+
+        # Call the method and assert the exception
+        with pytest.raises(
+            ValueError,
+            match="Please save the dataset before pushing to the remote server.",
+        ) as context:
+            sample_df.push()
+
+    @patch("pandasai.dataframe.base.os.environ")
+    def test_push_raises_error_if_api_key_is_missing(self, mock_environ, sample_df):
+        # Mock environment variable as missing
+        mock_environ.get.return_value = None
+
+        # Call the method and assert the exception
+        with pytest.raises(PandaAIApiKeyError):
+            sample_df.push()
+
+    @patch("pandasai.dataframe.base.os.path.exists")
+    @patch("pandasai.dataframe.base.open", new_callable=mock_open)
+    @patch("pandasai.dataframe.base.get_pandaai_session")
+    @patch("pandasai.dataframe.base.os.environ")
+    def test_push_closes_files_on_completion(
+        self,
+        mock_environ,
+        mock_get_session,
+        mock_open,
+        mock_path_exists,
+        sample_df: DataFrame,
+    ):
+        # Mock environment variable
+        mock_environ.get.return_value = "fake_api_key"
+        # Mock file existence
+        mock_path_exists.return_value = True
+
+        # Mock session
+        mock_session = MagicMock()
+        mock_get_session.return_value = mock_session
+
+        # Call the method
+        sample_df.push()
+
+        # Assert that files were closed after the request
+        mock_open.return_value.close.assert_called()

tests/unit_tests/dataframe/test_loader.py

@@ -377,3 +377,44 @@ def test_load_direct_from_sqlite_schema(self, sqlite_schema):
             assert mock_execute_query.call_args[0][0] == (
                 "SELECT email, first_name, timestamp FROM users ORDER BY created_at DESC LIMIT 100"
             )
+
+    def test_load_without_schema_and_path(self, sample_schema):
+        """Test load by not providing a schema or path."""
+        with patch("os.path.exists", return_value=True), patch(
+            "pandasai.data_loader.loader.DatasetLoader._read_csv_or_parquet"
+        ) as mock_read_csv_or_parquet, patch(
+            "pandasai.data_loader.loader.DatasetLoader._apply_transformations"
+        ) as mock_apply_transformations:
+            mock_read_csv_or_parquet.return_value = DataFrame(
+                {"email": ["[email protected]"]}
+            )
+            mock_apply_transformations.return_value = DataFrame(
+                {"email": ["[email protected]"]}
+            )
+            loader = DatasetLoader()
+
+            with pytest.raises(
+                ValueError, match="Either 'dataset_path' or 'schema' must be provided."
+            ):
+                result = loader.load()
+
+    def test_load_with_schema_and_path(self, sample_schema):
+        """Test load by providing both schema and path."""
+        with patch("os.path.exists", return_value=True), patch(
+            "pandasai.data_loader.loader.DatasetLoader._read_csv_or_parquet"
+        ) as mock_read_csv_or_parquet, patch(
+            "pandasai.data_loader.loader.DatasetLoader._apply_transformations"
+        ) as mock_apply_transformations:
+            mock_read_csv_or_parquet.return_value = DataFrame(
+                {"email": ["[email protected]"]}
+            )
+            mock_apply_transformations.return_value = DataFrame(
+                {"email": ["[email protected]"]}
+            )
+            loader = DatasetLoader()
+
+            with pytest.raises(
+                ValueError,
+                match="Provide only one of 'dataset_path' or 'schema', not both.",
+            ):
+                result = loader.load("test/users", sample_schema)