fix tests

Author: vvkh

data_diff/info_tree.py

@@ -1,4 +1,4 @@
-from typing import List, Dict, Optional, Any, Tuple
+from typing import List, Dict, Optional, Any, Tuple, Union
 
 from runtype import dataclass
 
@@ -9,15 +9,15 @@
 class SegmentInfo:
     tables: List[TableSegment]
 
-    diff: List[Tuple[Any, ...]] = None
+    diff: List[Union[Tuple[Any, ...], List[Any]]] = None
     diff_schema: Tuple[Tuple[str, type], ...] = None
     is_diff: bool = None
     diff_count: int = None
 
     rowcounts: Dict[int, int] = {}
     max_rows: int = None
 
-    def set_diff(self, diff: List[List[Any]], schema: Optional[Tuple[Tuple[str, type]]] = None):
+    def set_diff(self, diff: List[Union[Tuple[Any, ...], List[Any]]], schema: Optional[Tuple[Tuple[str, type]]] = None):
         self.diff_schema = schema
         self.diff = diff
         self.diff_count = len(diff)

tests/test_dbt.py

@@ -835,7 +835,7 @@ def test_diff_is_not_cloud(self, mock_print, mock_dbt_parser, mock_cloud_diff, m
         mock_dbt_parser_inst.get_models.assert_called_once()
         mock_dbt_parser_inst.set_connection.assert_called_once()
         mock_cloud_diff.assert_not_called()
-        mock_local_diff.assert_called_once_with(diff_vars)
+        mock_local_diff.assert_called_once_with(diff_vars, False)
         mock_print.assert_not_called()
 
     @patch("data_diff.dbt._get_diff_vars")
@@ -909,7 +909,7 @@ def test_diff_only_prod_db(self, mock_print, mock_dbt_parser, mock_cloud_diff, m
         mock_dbt_parser_inst.get_models.assert_called_once()
         mock_dbt_parser_inst.set_connection.assert_called_once()
         mock_cloud_diff.assert_not_called()
-        mock_local_diff.assert_called_once_with(diff_vars)
+        mock_local_diff.assert_called_once_with(diff_vars, False)
         mock_print.assert_not_called()
 
     @patch("data_diff.dbt._get_diff_vars")
@@ -946,7 +946,7 @@ def test_diff_only_prod_schema(
         mock_dbt_parser_inst.get_models.assert_called_once()
         mock_dbt_parser_inst.set_connection.assert_called_once()
         mock_cloud_diff.assert_not_called()
-        mock_local_diff.assert_called_once_with(diff_vars)
+        mock_local_diff.assert_called_once_with(diff_vars, False)
         mock_print.assert_not_called()
 
     @patch("data_diff.dbt._initialize_api")