dtenedor commented on code in PR #57691:
URL: https://github.com/apache/spark/pull/57691#discussion_r3716274852
##########
python/pyspark/sql/tests/df_golden/df_golden.py:
##########
@@ -744,3 +752,230 @@ def _compare_case(test_case, case, actual):
got.strip("\n"),
"[{}] mismatch in `{}`".format(name, key),
)
+
+
+def is_generating_golden():
+ """Whether this run regenerates the golden files instead of checking
them."""
+ return os.environ.get("SPARK_GENERATE_GOLDEN_FILES") is not None
+
+
+def check_cases_in_sync(test_file, declared, golden):
+ """
+ Check that the golden file describes exactly the declared cases, in order.
+
+ Each test asserts against its own block, so a case the golden file has
never
+ heard of (or one it still remembers after the method was deleted or
renamed)
+ would otherwise go unnoticed.
+ """
+ if declared == golden:
+ return
+ missing = [name for name in declared if name not in golden]
+ extra = [name for name in golden if name not in declared]
+ if missing:
+ detail = "cases with no block in the golden file: " + ",
".join(missing)
+ elif extra:
+ detail = "blocks in the golden file with no case method: " + ",
".join(extra)
+ else:
+ detail = "the golden file lists the cases in a different order"
+ raise AssertionError(
+ "{}: golden file is out of sync with the test class ({}); "
+ "regenerate the golden files".format(test_file, detail)
+ )
+
+
+# ---------------------------------------------------------------------------
+# Test class integration
+# ---------------------------------------------------------------------------
+
+
+class DFGoldenTestMixin:
+ """
+ Mixin turning a class of case methods into DataFrame golden file tests.
+
+ Mix into a session-providing test case, listing this class first so its
+ ``setUpClass`` runs once the session exists::
+
+ class GroupByGoldenTests(DFGoldenTestMixin, ReusedConnectTestCase):
+ golden_file = "group_by.test"
+
+ def _test_group_by_count(self, spark):
+ return
spark.table("testData").groupBy(col("a")).agg(count(col("b")))
+
+ Every ``_test_<case>`` method declares one case and returns the DataFrame
+ under test; a ``test_<case>`` method is registered for each, so cases run,
+ report and can be selected individually like any other unittest test.
+
+ The cases of a class share one Spark Connect session (``newSession()`` off
+ the session the test class provides, the Connect counterpart of
+ ``SQLQueryTestSuite``'s per-file ``newSession()``), prepared once by
+ :meth:`setup_session`. State created there -- temp views, UDFs, session
+ confs -- is discarded with the session and cannot leak into other classes.
+ """
+
+ #: Name of the ``.test`` golden file, resolved next to the test module.
+ golden_file = None
+
+ def __init_subclass__(cls, **kwargs):
+ super().__init_subclass__(**kwargs)
+ # This mixin's setUpClass must run after the session-providing class
has
+ # created the session, i.e. its setUpClass must be the outer one, which
+ # is only true when the mixin is listed first.
+ for base in cls.__mro__:
+ if base is DFGoldenTestMixin:
+ break
+ if base is not cls and "setUpClass" in vars(base):
+ raise TypeError(
+ "{} has incorrect inheritance order: DFGoldenTestMixin
must be "
+ "listed before {}. Use: class {}(DFGoldenTestMixin, {},
...)".format(
+ cls.__name__, base.__name__, cls.__name__,
base.__name__
+ )
+ )
+ for name in cls.case_names():
+ setattr(cls, "test_" + name, _make_case_test(name,
cls.case_method(name)))
+
+ @classmethod
+ def case_names(cls):
+ """The declared case names, in declaration order."""
+ names = []
+ for klass in reversed(cls.__mro__):
+ for attr, value in vars(klass).items():
+ if not attr.startswith(_CASE_METHOD_PREFIX) or not
callable(value):
+ continue
+ name = attr[len(_CASE_METHOD_PREFIX) :]
+ if name not in names:
+ names.append(name)
+ return names
+
+ @classmethod
+ def case_method(cls, name):
+ """The case method declaring the case *name*."""
+ return getattr(cls, _CASE_METHOD_PREFIX + name)
+
+ @classmethod
+ def golden_file_path(cls):
+ """Absolute path of the golden file, resolved next to the test
module."""
+ assert cls.golden_file, "{}: set `golden_file` to its .test
file".format(cls.__name__)
+ module_file = os.path.abspath(inspect.getfile(cls))
+ return os.path.join(os.path.dirname(module_file), cls.golden_file)
+
+ @classmethod
+ def setup_session(cls, spark):
+ """
+ Hook: prepare the session shared by this class's cases.
+
+ Override to create the temp views and other session state the cases
+ build on.
+ """
+
+ @classmethod
+ def setUpClass(cls):
+ cls._golden_regenerating = is_generating_golden()
+ cls._golden_actual = {}
+ cls._golden_cases = {}
+ cls._golden_session = None
+ case_names = cls.case_names()
+
+ # Read the golden file before the session is started: a malformed or
+ # stale file is worth failing on right away, and failing here leaves
+ # nothing to clean up (unittest skips tearDownClass when setUpClass
+ # raises). Regeneration writes the file from the class alone, so it
+ # neither reads nor requires an existing one.
+ if case_names:
+ cls._golden_path = cls.golden_file_path()
+ if not cls._golden_regenerating:
+ header, cases = parse_test_file(cls._golden_path)
+ validate_test_file(cls._golden_path, header, cases)
+ golden_names = [case["name"] for case in cases]
+ check_cases_in_sync(cls._golden_path, case_names, golden_names)
Review Comment:
This is done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
