Add --code flag for collecting code context. (#1154)

* Added --code flag for collecting code context. * Addressed reviewer request and fixed tests * Fixed decoration problems and added tests * Addressed review requests * Refactored solution by review requests * Uncommented line in LogfileExpect --------- Co-authored-by: Robert Nemes <[email protected]> Co-authored-by: Pyifan <[email protected]>
morganstanley · Dec 19, 2024 · 9f93423 · 9f93423
1 parent d9520ee
commit 9f93423
Show file tree

Hide file tree

Showing 18 changed files with 280 additions and 77 deletions.
diff --git a/doc/en/introduction.rst b/doc/en/introduction.rst
@@ -455,6 +455,7 @@ Command line
 
       --label LABEL         Label the test report with the given name, useful to categorize or classify similar reports (aka "run-id").
       --driver-info         Display drivers startup and teardown information, and visualise driver connections in the report.
+      --code                Collects file path, line number and code context of the assertions.
 
 
 Highlighted features

diff --git a/doc/newsfragments/2981_new.code_context.rst b/doc/newsfragments/2981_new.code_context.rst
@@ -0,0 +1,2 @@
+Added ``--code`` flag to collect code context for the assertions. Code context one-liner will be displayed on the web UI if enabled.
+Note that file path information is no longer collected by default. To collect file path information, enable code context.
diff --git a/testplan/base.py b/testplan/base.py
@@ -146,6 +146,8 @@ class Testplan(entity.RunnableManager):
         categorize or classify similar reports .
     :param driver_info: Display driver setup / teardown time and driver
         interconnection information in UI report.
+    :param collect_code_context: Collects the file path, line number and code
+        context of the assertions.
     """
 
     CONFIG = TestplanConfig
@@ -194,6 +196,7 @@ def __init__(
         extra_deps: Optional[List[Union[str, ModuleType]]] = None,
         label: Optional[str] = None,
         driver_info: bool = False,
+        collect_code_context: bool = False,
         auto_part_runtime_limit: int = defaults.AUTO_PART_RUNTIME_LIMIT,
         plan_runtime_target: int = defaults.PLAN_RUNTIME_TARGET,
         **options,
@@ -256,6 +259,7 @@ def __init__(
             extra_deps=extra_deps,
             label=label,
             driver_info=driver_info,
+            collect_code_context=collect_code_context,
             auto_part_runtime_limit=auto_part_runtime_limit,
             plan_runtime_target=plan_runtime_target,
             **options,
@@ -401,6 +405,7 @@ def main_wrapper(
         extra_deps=None,
         label=None,
         driver_info=False,
+        collect_code_context=False,
         auto_part_runtime_limit=defaults.AUTO_PART_RUNTIME_LIMIT,
         plan_runtime_target=defaults.PLAN_RUNTIME_TARGET,
         **options,
@@ -462,6 +467,7 @@ def test_plan_inner_inner():
                     extra_deps=extra_deps,
                     label=label,
                     driver_info=driver_info,
+                    collect_code_context=collect_code_context,
                     auto_part_runtime_limit=auto_part_runtime_limit,
                     plan_runtime_target=plan_runtime_target,
                     **options,

diff --git a/testplan/parser.py b/testplan/parser.py
@@ -499,6 +499,14 @@ def generate_parser(self) -> HelpParser:
             help="Display drivers setup / teardown timing and interconnection information in UI report.",
         )
 
+        report_group.add_argument(
+            "--code",
+            dest="collect_code_context",
+            action="store_true",
+            default=self._default_options["collect_code_context"],
+            help="Collects file path, line number and code context of the assertions.",
+        )
+
         self.add_arguments(parser)
         return parser
 

diff --git a/testplan/runnable/base.py b/testplan/runnable/base.py
@@ -267,6 +267,7 @@ def get_options(cls):
                 "skip_strategy", default=common.SkipStrategy.noop()
             ): Use(common.SkipStrategy.from_option_or_none),
             ConfigOption("driver_info", default=False): bool,
+            ConfigOption("collect_code_context", default=False): bool,
         }
 
 

diff --git a/testplan/testing/base.py b/testplan/testing/base.py
@@ -693,7 +693,9 @@ def _run_resource_hook(
         )
 
         case_result = self.cfg.result(
-            stdout_style=self.stdout_style, _scratch=self.scratch
+            stdout_style=self.stdout_style,
+            _scratch=self.scratch,
+            _collect_code_context=self.collect_code_context,
         )
         runtime_env = self._get_runtime_environment(
             testcase_name=hook_name,
@@ -901,6 +903,14 @@ def driver_info(self) -> bool:
             return False
         return self.cfg.driver_info
 
+    @property
+    def collect_code_context(self) -> bool:
+        """
+        Collecting the file path, line number and code context of the assertions
+        if enabled.
+        """
+        return getattr(self.cfg, "collect_code_context", False)
+
 
 class ProcessRunnerTestConfig(TestConfig):
     """

diff --git a/testplan/testing/multitest/base.py b/testplan/testing/multitest/base.py
@@ -1045,7 +1045,9 @@ def _run_suite_related(self, testsuite, method_name):
 
         method_report = self._suite_related_report(method_name)
         case_result = self.cfg.result(
-            stdout_style=self.stdout_style, _scratch=self._scratch
+            stdout_style=self.stdout_style,
+            _scratch=self._scratch,
+            _collect_code_context=self.collect_code_context,
         )
 
         resources = self._get_runtime_environment(
@@ -1137,7 +1139,9 @@ def _run_testcase(
             testcase
         )
         case_result: result.Result = self.cfg.result(
-            stdout_style=self.stdout_style, _scratch=self.scratch
+            stdout_style=self.stdout_style,
+            _scratch=self.scratch,
+            _collect_code_context=self.collect_code_context,
         )
 
         # as the runtime info currently has only testcase name we create it here
@@ -1148,7 +1152,7 @@ def _run_testcase(
             testcase_name=testcase.name, testcase_report=testcase_report
         )
 
-        if self.cfg.testcase_report_target:
+        if self.cfg.testcase_report_target and self.collect_code_context:
             testcase = report_target(
                 func=testcase,
                 ref_func=getattr(

diff --git a/testplan/testing/multitest/entries/base.py b/testplan/testing/multitest/entries/base.py
@@ -50,6 +50,7 @@ def __init__(self, description, category=None, flag=None):
         # Will be set explicitly via containers
         self.line_no = None
         self.file_path = None
+        self.code_context = None
 
     def __str__(self):
         return repr(self)

diff --git a/testplan/testing/multitest/entries/schemas/base.py b/testplan/testing/multitest/entries/schemas/base.py
@@ -33,6 +33,7 @@ class BaseSchema(Schema):
     category = fields.String()
     flag = fields.String()
     file_path = fields.String()
+    code_context = fields.String()
     custom_style = fields.Dict(keys=fields.String(), values=fields.String())
 
     def load(self, *args, **kwargs):

diff --git a/testplan/testing/result.py b/testplan/testing/result.py
@@ -93,12 +93,14 @@ def __exit__(self, exc_type, exc_value, tb):
             description=self.description,
         )
 
-        with MOD_LOCK:
-            # TODO: see https://github.com/python/cpython/commit/85cf1d514b84dc9a4bcb40e20a12e1d82ff19f20
-            caller_frame = inspect.stack()[1]
+        if self.result._collect_code_context:
+            with MOD_LOCK:
+                # TODO: see https://github.com/python/cpython/commit/85cf1d514b84dc9a4bcb40e20a12e1d82ff19f20
+                caller_frame = inspect.stack()[1]
 
-        exc_assertion.file_path = os.path.abspath(caller_frame[1])
-        exc_assertion.line_no = caller_frame[2]
+            exc_assertion.file_path = os.path.abspath(caller_frame[1])
+            exc_assertion.line_no = caller_frame[2]
+            exc_assertion.code_context = caller_frame.code_context[0].strip()
 
         # We cannot use `bind_entry` here as this block will
         # be run when an exception is raised
@@ -165,7 +167,24 @@ def wrapper(result, *args, **kwargs):
             custom_style = kwargs.pop("custom_style", None)
             dryrun = kwargs.pop("dryrun", False)
             entry = func(result, *args, **kwargs)
-            if top_assertion:
+            if not top_assertion:
+                return entry
+
+            if custom_style is not None:
+                if not isinstance(custom_style, dict):
+                    raise TypeError(
+                        "Use `dict[str, str]` to specify custom CSS style"
+                    )
+                entry.custom_style = custom_style
+
+            assert isinstance(result, AssertionNamespace) or isinstance(
+                result, Result
+            ), "Incorrect usage of assertion decorator"
+
+            if isinstance(result, AssertionNamespace):
+                result = result.result
+
+            if result._collect_code_context:
                 with MOD_LOCK:
                     call_stack = inspect.stack()
                     try:
@@ -183,34 +202,21 @@ def wrapper(result, *args, **kwargs):
                                 frame = call_stack[1]
                         entry.file_path = os.path.abspath(frame.filename)
                         entry.line_no = frame.lineno
+                        entry.code_context = frame.code_context[0].strip()
                     finally:
                         # https://docs.python.org/3/library/inspect.html
                         del frame
                         del call_stack
 
-                if custom_style is not None:
-                    if not isinstance(custom_style, dict):
-                        raise TypeError(
-                            "Use `dict[str, str]` to specify custom CSS style"
-                        )
-                    entry.custom_style = custom_style
-
-                assert isinstance(result, AssertionNamespace) or isinstance(
-                    result, Result
-                ), "Incorrect usage of assertion decorator"
+            if not dryrun:
+                result.entries.append(entry)
 
-                if isinstance(result, AssertionNamespace):
-                    result = result.result
-
-                if not dryrun:
-                    result.entries.append(entry)
-
-                stdout_registry.log_entry(
-                    entry=entry, stdout_style=result.stdout_style
-                )
+            stdout_registry.log_entry(
+                entry=entry, stdout_style=result.stdout_style
+            )
 
-                if not entry and not result.continue_on_failure:
-                    raise AssertionError(entry)
+            if not entry and not result.continue_on_failure:
+                raise AssertionError(entry)
 
             return entry
         finally:
@@ -1371,10 +1377,13 @@ def __exit__(self, exc_type, exc_value, traceback):
             return False
         super().__exit__(exc_type, exc_value, traceback)
 
-        with MOD_LOCK:
-            # TODO: see https://github.com/python/cpython/commit/85cf1d514b84dc9a4bcb40e20a12e1d82ff19f20
-            # XXX: do we have concrete ideas about thread-safety here?
-            caller_frame = inspect.stack()[1]
+        if self.result._collect_code_context:
+            with MOD_LOCK:
+                # TODO: see https://github.com/python/cpython/commit/85cf1d514b84dc9a4bcb40e20a12e1d82ff19f20
+                # XXX: do we have concrete ideas about thread-safety here?
+                caller_frame = inspect.stack()[1]
+        else:
+            caller_frame = None
 
         assertion = assertions.LogfileMatch(
             self.timeout,
@@ -1383,8 +1392,11 @@ def __exit__(self, exc_type, exc_value, traceback):
             self.description,
             self.category,
         )
-        assertion.file_path = os.path.abspath(caller_frame[1])
-        assertion.line_no = caller_frame[2]
+
+        if caller_frame:
+            assertion.file_path = os.path.abspath(caller_frame[1])
+            assertion.line_no = caller_frame[2]
+            assertion.code_context = caller_frame.code_context[0].strip()
 
         stdout_registry.log_entry(
             entry=assertion, stdout_style=self.result.stdout_style
@@ -1534,6 +1546,7 @@ def __init__(
         _num_passing=defaults.SUMMARY_NUM_PASSING,
         _num_failing=defaults.SUMMARY_NUM_FAILING,
         _scratch=None,
+        _collect_code_context=False,
     ):
 
         self.entries = []
@@ -1555,6 +1568,7 @@ def __init__(
         self._num_passing = _num_passing
         self._num_failing = _num_failing
         self._scratch = _scratch
+        self._collect_code_context = _collect_code_context
 
     def subresult(self):
         """Subresult object to append/prepend assertions on another."""
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		Added ``--code`` flag to collect code context for the assertions. Code context one-liner will be displayed on the web UI if enabled.
		Note that file path information is no longer collected by default. To collect file path information, enable code context.