1 files changed, 2 insertions, 3 deletions
diff --git a/tests/dataframe/unit/dataframe_sql_validator.py b/tests/dataframe/unit/dataframe_sql_validator.py
index 32ff8f2..2dcdb39 100644
--- a/tests/dataframe/unit/dataframe_sql_validator.py
+++ b/tests/dataframe/unit/dataframe_sql_validator.py
@@ -4,6 +4,7 @@ import unittest
 from sqlglot.dataframe.sql import types
 from sqlglot.dataframe.sql.dataframe import DataFrame
 from sqlglot.dataframe.sql.session import SparkSession
+from sqlglot.helper import ensure_list
 
 
 class DataFrameSQLValidator(unittest.TestCase):
@@ -33,9 +34,7 @@ class DataFrameSQLValidator(unittest.TestCase):
         self, df: DataFrame, expected_statements: t.Union[str, t.List[str]], pretty=False
     ):
         actual_sqls = df.sql(pretty=pretty)
-        expected_statements = (
-            [expected_statements] if isinstance(expected_statements, str) else expected_statements
-        )
+        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)