gaogaotiantian commented on code in PR #57691:
URL: https://github.com/apache/spark/pull/57691#discussion_r3788173913
##########
python/pyspark/sql/tests/df_golden/df_golden.py:
##########
@@ -744,3 +752,225 @@ 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 golden_file_for_input(module_file):
+ """
+ Return the golden file path for the test module at *module_file*.
+
+ Mirrors ``SQLQueryTestSuite.resultFileForInputFile``: the ``inputs``
+ directory becomes ``results`` and ``.out`` is appended, so a module and its
+ golden file are named alike and there is nothing to keep in sync by hand.
+ """
+ inputs_dir, filename = os.path.split(os.path.abspath(module_file))
+ base_dir, dir_name = os.path.split(inputs_dir)
+ assert dir_name == _INPUTS_DIR, (
+ "{}: a golden test module must live in the `{}` directory, next to the
"
+ "`{}` directory holding the golden files".format(module_file,
_INPUTS_DIR, _RESULTS_DIR)
+ )
+ return os.path.join(base_dir, _RESULTS_DIR, filename + ".out")
+
+
+def assert_cases_in_sync(test_file, declared, golden):
+ """
+ Assert 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):
+ 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
+ golden file is derived from the test module (see
+ :func:`golden_file_for_input`), so there is nothing to declare.
+
+ 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.
+ """
+
+ def __init_subclass__(cls, **kwargs):
+ super().__init_subclass__(**kwargs)
+ # Only the cases this class declares: an inherited one already has its
+ # test method on the class that declares it.
+ for attr, method in list(vars(cls).items()):
+ if attr.startswith(_CASE_METHOD_PREFIX) and
inspect.isfunction(method):
+ cls._register_case_test(attr[len(_CASE_METHOD_PREFIX) :],
method)
+
+ @classmethod
+ def _register_case_test(cls, name, case_method):
+ """Add the unittest method that runs the case *name*."""
+
+ def test_case(self):
+ self._run_golden_case(name)
+
+ test_case.__name__ = "test_" + name
+ # Carry the case method's docstring over so verbose runs describe the
case.
+ test_case.__doc__ = case_method.__doc__
+ setattr(cls, test_case.__name__, test_case)
+
+ @classmethod
+ def case_names(cls):
+ """The declared case names, including inherited ones, sorted by
name."""
+ return [
+ name[len(_CASE_METHOD_PREFIX) :]
+ for name, _ in inspect.getmembers(cls,
predicate=inspect.isfunction)
+ if name.startswith(_CASE_METHOD_PREFIX)
+ ]
+
+ @classmethod
+ def golden_file_path(cls):
+ """Absolute path of this class's golden file, derived from its
module."""
+ return golden_file_for_input(inspect.getfile(cls))
+
+ @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.
+ """
+ pass
+
+ @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:
+ cls._golden_cases = parse_test_file(cls._golden_path)
+ assert_cases_in_sync(cls._golden_path, case_names,
list(cls._golden_cases))
+
+ super().setUpClass()
+
+ if not case_names:
+ return
+ # This mixin has to precede the session-providing test class in the
base
+ # list, so that its setUpClass body runs once that class has created
the
+ # session. Checking here rather than when the class is declared keeps
+ # legitimate chains -- a subclass with a setUpClass of its own --
working.
+ assert "spark" in vars(cls), (
+ "{}: no session was created; list DFGoldenTestMixin before the "
+ "session-providing test class".format(cls.__name__)
+ )
+ try:
+ cls._golden_session = cls.spark.newSession()
+ # Golden files are generated with ANSI mode on, matching the SQL
+ # golden tests. The session is discarded with the class, so there
+ # is nothing to restore.
+ cls._golden_session.conf.set("spark.sql.ansi.enabled", "true")
Review Comment:
If the server is shutdown, we don't need to release the session anymore
right? Can we do something smart in `_close_golden_session` to deal with this?
It would be nice if we can isolate all the release logic in a single method and
don't need to worry about the order of release.
--
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]
