Revert "remove changes in filters tests"

This reverts commit f3f795a.

Author: varun-edachali-dbx

tests/unit/test_filters.py

@@ -4,6 +4,11 @@
 
 import unittest
 from unittest.mock import MagicMock, patch
+import sys
+from typing import List, Dict, Any
+
+# Add the necessary path to import the filter module
+sys.path.append("/home/varun.edachali/conn/databricks-sql-python/src")
 
 from databricks.sql.backend.filters import ResultSetFilter
 
@@ -15,39 +20,24 @@ def setUp(self):
         """Set up test fixtures."""
         # Create a mock SeaResultSet
         self.mock_sea_result_set = MagicMock()
-
-        # Set up the remaining_rows method on the results attribute
-        self.mock_sea_result_set.results = MagicMock()
-        self.mock_sea_result_set.results.remaining_rows.return_value = [
-            ["catalog1", "schema1", "table1", "owner1", "2023-01-01", "TABLE", ""],
-            ["catalog1", "schema1", "table2", "owner1", "2023-01-01", "VIEW", ""],
-            [
-                "catalog1",
-                "schema1",
-                "table3",
-                "owner1",
-                "2023-01-01",
-                "SYSTEM TABLE",
-                "",
-            ],
-            [
-                "catalog1",
-                "schema1",
-                "table4",
-                "owner1",
-                "2023-01-01",
-                "EXTERNAL TABLE",
-                "",
-            ],
-        ]
+        self.mock_sea_result_set._response = {
+            "result": {
+                "data_array": [
+                    ["catalog1", "schema1", "table1", "TABLE", ""],
+                    ["catalog1", "schema1", "table2", "VIEW", ""],
+                    ["catalog1", "schema1", "table3", "SYSTEM TABLE", ""],
+                    ["catalog1", "schema1", "table4", "EXTERNAL TABLE", ""],
+                ],
+                "row_count": 4,
+            }
+        }
 
         # Set up the connection and other required attributes
         self.mock_sea_result_set.connection = MagicMock()
         self.mock_sea_result_set.backend = MagicMock()
         self.mock_sea_result_set.buffer_size_bytes = 1000
         self.mock_sea_result_set.arraysize = 100
         self.mock_sea_result_set.statement_id = "test-statement-id"
-        self.mock_sea_result_set.lz4_compressed = False
 
         # Create a mock CommandId
         from databricks.sql.backend.types import CommandId, BackendType
@@ -60,102 +50,70 @@ def setUp(self):
             ("catalog_name", "string", None, None, None, None, True),
             ("schema_name", "string", None, None, None, None, True),
             ("table_name", "string", None, None, None, None, True),
-            ("owner", "string", None, None, None, None, True),
-            ("creation_time", "string", None, None, None, None, True),
             ("table_type", "string", None, None, None, None, True),
             ("remarks", "string", None, None, None, None, True),
         ]
         self.mock_sea_result_set.has_been_closed_server_side = False
-        self.mock_sea_result_set._arrow_schema_bytes = None
 
-    def test_filter_by_column_values(self):
-        """Test filtering by column values with various options."""
-        # Case 1: Case-sensitive filtering
-        allowed_values = ["table1", "table3"]
+    def test_filter_tables_by_type(self):
+        """Test filtering tables by type."""
+        # Test with specific table types
+        table_types = ["TABLE", "VIEW"]
 
+        # Make the mock_sea_result_set appear to be a SeaResultSet
         with patch("databricks.sql.backend.filters.isinstance", return_value=True):
             with patch(
                 "databricks.sql.result_set.SeaResultSet"
             ) as mock_sea_result_set_class:
+                # Set up the mock to return a new mock when instantiated
                 mock_instance = MagicMock()
                 mock_sea_result_set_class.return_value = mock_instance
 
-                # Call filter_by_column_values on the table_name column (index 2)
-                result = ResultSetFilter.filter_by_column_values(
-                    self.mock_sea_result_set, 2, allowed_values, case_sensitive=True
+                result = ResultSetFilter.filter_tables_by_type(
+                    self.mock_sea_result_set, table_types
                 )
 
                 # Verify the filter was applied correctly
                 mock_sea_result_set_class.assert_called_once()
 
-                # Check the filtered data passed to the constructor
-                args, kwargs = mock_sea_result_set_class.call_args
-                result_data = kwargs.get("result_data")
-                self.assertIsNotNone(result_data)
-                self.assertEqual(len(result_data.data), 2)
-                self.assertIn(result_data.data[0][2], allowed_values)
-                self.assertIn(result_data.data[1][2], allowed_values)
+    def test_filter_tables_by_type_case_insensitive(self):
+        """Test filtering tables by type with case insensitivity."""
+        # Test with lowercase table types
+        table_types = ["table", "view"]
 
-        # Case 2: Case-insensitive filtering
-        mock_sea_result_set_class.reset_mock()
+        # Make the mock_sea_result_set appear to be a SeaResultSet
         with patch("databricks.sql.backend.filters.isinstance", return_value=True):
             with patch(
                 "databricks.sql.result_set.SeaResultSet"
             ) as mock_sea_result_set_class:
+                # Set up the mock to return a new mock when instantiated
                 mock_instance = MagicMock()
                 mock_sea_result_set_class.return_value = mock_instance
 
-                # Call filter_by_column_values with case-insensitive matching
-                result = ResultSetFilter.filter_by_column_values(
-                    self.mock_sea_result_set,
-                    2,
-                    ["TABLE1", "TABLE3"],
-                    case_sensitive=False,
-                )
-                mock_sea_result_set_class.assert_called_once()
-
-        # Case 3: Unsupported result set type
-        mock_unsupported_result_set = MagicMock()
-        with patch("databricks.sql.backend.filters.isinstance", return_value=False):
-            with patch("databricks.sql.backend.filters.logger") as mock_logger:
-                result = ResultSetFilter.filter_by_column_values(
-                    mock_unsupported_result_set, 0, ["value"], True
+                result = ResultSetFilter.filter_tables_by_type(
+                    self.mock_sea_result_set, table_types
                 )
-                mock_logger.warning.assert_called_once()
-                self.assertEqual(result, mock_unsupported_result_set)
 
-    def test_filter_tables_by_type(self):
-        """Test filtering tables by type with various options."""
-        # Case 1: Specific table types
-        table_types = ["TABLE", "VIEW"]
+                # Verify the filter was applied correctly
+                mock_sea_result_set_class.assert_called_once()
 
+    def test_filter_tables_by_type_default(self):
+        """Test filtering tables by type with default types."""
+        # Make the mock_sea_result_set appear to be a SeaResultSet
         with patch("databricks.sql.backend.filters.isinstance", return_value=True):
-            with patch.object(
-                ResultSetFilter, "filter_by_column_values"
-            ) as mock_filter:
-                ResultSetFilter.filter_tables_by_type(
-                    self.mock_sea_result_set, table_types
+            with patch(
+                "databricks.sql.result_set.SeaResultSet"
+            ) as mock_sea_result_set_class:
+                # Set up the mock to return a new mock when instantiated
+                mock_instance = MagicMock()
+                mock_sea_result_set_class.return_value = mock_instance
+
+                result = ResultSetFilter.filter_tables_by_type(
+                    self.mock_sea_result_set, None
                 )
-                args, kwargs = mock_filter.call_args
-                self.assertEqual(args[0], self.mock_sea_result_set)
-                self.assertEqual(args[1], 5)  # Table type column index
-                self.assertEqual(args[2], table_types)
-                self.assertEqual(kwargs.get("case_sensitive"), True)
 
-        # Case 2: Default table types (None or empty list)
-        with patch("databricks.sql.backend.filters.isinstance", return_value=True):
-            with patch.object(
-                ResultSetFilter, "filter_by_column_values"
-            ) as mock_filter:
-                # Test with None
-                ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, None)
-                args, kwargs = mock_filter.call_args
-                self.assertEqual(args[2], ["TABLE", "VIEW", "SYSTEM TABLE"])
-
-                # Test with empty list
-                ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, [])
-                args, kwargs = mock_filter.call_args
-                self.assertEqual(args[2], ["TABLE", "VIEW", "SYSTEM TABLE"])
+                # Verify the filter was applied correctly
+                mock_sea_result_set_class.assert_called_once()
 
 
 if __name__ == "__main__":