[library/python/pytest/yatest_tools] Fixed module discovering for tests with CONFTEST_LOAD_POLICY_LOCAL()

1 files changed, 70 insertions, 17 deletions

diff --git a/library/python/pytest/yatest_tools.py b/library/python/pytest/yatest_tools.py
index 62042a2bdc..723be996ed 100644
--- a/library/python/pytest/yatest_tools.py
+++ b/library/python/pytest/yatest_tools.py
@@ -7,9 +7,12 @@ import os
 import re
 import sys
 
+from . import config
 import yatest_lib.tools
 
-SEP = "/"
+
+SEP = '/'
+TEST_MOD_PREFIX = '__tests__.'
 
 
 class Subtest(object):
@@ -155,13 +158,19 @@ TRACE_FILE_NAME = "ytest.report.trace"
 
 
 def lazy(func):
-    mem = {}
+    memory = {}
 
     @functools.wraps(func)
-    def wrapper():
-        if "results" not in mem:
-            mem["results"] = func()
-        return mem["results"]
+    def wrapper(*args):
+        # Disabling caching in test mode
+        if config.is_test_mode():
+            return func(*args)
+
+        try:
+            return memory[args]
+        except KeyError:
+            memory[args] = func(*args)
+        return memory[args]
 
     return wrapper
 
@@ -175,6 +184,7 @@ def _get_mtab():
     return []
 
 
+@lazy
 def get_max_filename_length(dirname):
     """
     Return maximum filename length for the filesystem
@@ -253,6 +263,7 @@ def normalize_name(name):
     return name
 
 
+@lazy
 def normalize_filename(filename):
     """
     Replace invalid for file names characters with string equivalents
@@ -285,6 +296,7 @@ def get_test_log_file_path(output_dir, class_name, test_name, extension="log"):
     return get_unique_file_path(output_dir, filename)
 
 
+@lazy
 def split_node_id(nodeid, test_suffix=None):
     path, possible_open_bracket, params = nodeid.partition('[')
     separator = "::"
@@ -308,25 +320,66 @@ def split_node_id(nodeid, test_suffix=None):
     return yatest_lib.tools.to_utf8(class_name), yatest_lib.tools.to_utf8(test_name)
 
 
+@lazy
+def _suffix_test_modules_tree():
+    root = {}
+
+    for module in sys.extra_modules:
+        if not module.startswith(TEST_MOD_PREFIX):
+            continue
+
+        module = module[len(TEST_MOD_PREFIX):]
+        node = root
+
+        for name in reversed(module.split('.')):
+            if name == '__init__':
+                continue
+            node = node.setdefault(name, {})
+
+    return root
+
+
+def _conftest_load_policy_is_local(path):
+    return SEP in path and getattr(sys, "is_standalone_binary", False)
+
+
+class MissingTestModule(Exception):
+    pass
+
+
 # If CONFTEST_LOAD_POLICY==LOCAL the path parameters is a true test file path. Something like
 #   /-B/taxi/uservices/services/alt/gen/tests/build/services/alt/validation/test_generated_files.py
 # If CONFTEST_LOAD_POLICY is not LOCAL the path parameter is a module name with '.py' extension added. Example:
 #  validation.test_generated_files.py
 # To make test names independent of the CONFTEST_LOAD_POLICY value replace path by module name if possible.
+@lazy
 def _unify_path(path):
-    test_mod_pfx = "__tests__."
     py_ext = ".py"
 
     path = path.strip()
-    if SEP in path and getattr(sys, "is_standalone_binary", False):
-        # Try to find path as a module in test modules and use it as a class name
+    if _conftest_load_policy_is_local(path) and path.endswith(py_ext):
+        # Try to find best match for path as a module among test modules and use it as a class name.
         # This is the only way to unify different CONFTEST_LOAD_POLICY modes
+        suff_tree = _suffix_test_modules_tree()
+        node, res = suff_tree, []
+
+        assert path.endswith(py_ext), path
         parts = path[:-len(py_ext)].split(SEP)
-        pattern = "." + ".".join(parts)
-        for module in sys.extra_modules:
-            if module.startswith(test_mod_pfx):
-                m = module[len(test_mod_pfx) - 1:]
-                if pattern.endswith(m):
-                    # Remove leading '.' and add file extension
-                    return m[1:] + py_ext
-    return path
+
+        for p in reversed(parts):
+            if p in node:
+                node = node[p]
+                res.append(p)
+            else:
+                if res:
+                    return '.'.join(reversed(res)) + py_ext
+                else:
+                    # Top level test module
+                    if TEST_MOD_PREFIX + p in sys.extra_modules:
+                        return p + py_ext
+                    # Unknown module - raise an error
+                    break
+
+        raise MissingTestModule("Can't find proper module for '{}' path among: {}".format(path, suff_tree))
+    else:
+        return path