From f73e9af131151f1e058446361c35b05c4c90bf10 Mon Sep 17 00:00:00 2001 From: Daniel Baumann Date: Thu, 7 Sep 2023 13:39:48 +0200 Subject: Merging upstream version 18.2.0. Signed-off-by: Daniel Baumann --- tests/dataframe/unit/dataframe_test_base.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 tests/dataframe/unit/dataframe_test_base.py (limited to 'tests/dataframe/unit/dataframe_test_base.py') diff --git a/tests/dataframe/unit/dataframe_test_base.py b/tests/dataframe/unit/dataframe_test_base.py new file mode 100644 index 0000000..6b07df9 --- /dev/null +++ b/tests/dataframe/unit/dataframe_test_base.py @@ -0,0 +1,23 @@ +import typing as t +import unittest + +import sqlglot +from sqlglot import MappingSchema +from sqlglot.dataframe.sql import SparkSession +from sqlglot.dataframe.sql.dataframe import DataFrame +from sqlglot.helper import ensure_list + + +class DataFrameTestBase(unittest.TestCase): + def setUp(self) -> None: + sqlglot.schema = MappingSchema() + SparkSession._instance = None + + def compare_sql( + self, df: DataFrame, expected_statements: t.Union[str, t.List[str]], pretty=False + ): + actual_sqls = df.sql(pretty=pretty) + expected_statements = ensure_list(expected_statements) + self.assertEqual(len(expected_statements), len(actual_sqls)) + for expected, actual in zip(expected_statements, actual_sqls): + self.assertEqual(expected, actual) -- cgit v1.2.3