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

3 files changed, 116 insertions, 24 deletions

diff --git a/library/python/pytest/config.py b/library/python/pytest/config.py
new file mode 100644
index 0000000000..703e442c0b
--- /dev/null
+++ b/library/python/pytest/config.py
@@ -0,0 +1,9 @@
+_test_mode = [False]
+
+
+def is_test_mode():
+    return _test_mode[0]
+
+
+def set_test_mode():
+    _test_mode[0] = True
diff --git a/library/python/pytest/ut/test_tools.py b/library/python/pytest/ut/test_tools.py
index 6368b9c2fa..de85b1a6e0 100644
--- a/library/python/pytest/ut/test_tools.py
+++ b/library/python/pytest/ut/test_tools.py
@@ -1,7 +1,11 @@
 import pytest
 import sys
 
-from library.python.pytest.yatest_tools import split_node_id
+from library.python.pytest import config
+from library.python.pytest import yatest_tools
+
+
+config.set_test_mode()
 
 
 @pytest.fixture(params=["", "[param1,param2]"])
@@ -23,7 +27,7 @@ def parameters(request):
     ),
 )
 def test_split_node_id_without_path(parameters, node_id, expected_class_name, expected_test_name):
-    got = split_node_id(node_id + parameters)
+    got = yatest_tools.split_node_id(node_id + parameters)
     assert (expected_class_name, expected_test_name + parameters) == got
 
 
@@ -42,16 +46,20 @@ def test_split_node_id_without_path(parameters, node_id, expected_class_name, ex
             "package.test_script.py::class_name",
             "test_name",
         ),
-        # If module is not found in sys.extra_modules use basename as a class name
-        ("/arcadia/root/package/test_script2.py::test_name", "test_script2.py", "test_name"),
     ),
 )
 def test_split_node_id_with_path(mocker, parameters, node_id, expected_class_name, expected_test_name):
     mocker.patch.object(sys, 'extra_modules', sys.extra_modules | {'__tests__.package.test_script'})
-    got = split_node_id(node_id + parameters)
+    got = yatest_tools.split_node_id(node_id + parameters)
     assert (expected_class_name, expected_test_name + parameters) == got
 
 
+def test_missing_module(parameters):
+    # An errno must be raised if module is not found in sys.extra_modules
+    with pytest.raises(yatest_tools.MissingTestModule):
+        yatest_tools.split_node_id("/arcadia/root/package/test_script2.py::test_name" + parameters)
+
+
 @pytest.mark.parametrize(
     "node_id,expected_class_name,expected_test_name",
     (
@@ -61,6 +69,28 @@ def test_split_node_id_with_path(mocker, parameters, node_id, expected_class_nam
         ("package.test_script.py::class_name::subclass_name::test_name", "package.test_script.py", "test_suffix"),
     ),
 )
-def test_split_node_id_with_test_suffix(mocker, parameters, node_id, expected_class_name, expected_test_name):
-    got = split_node_id(node_id + parameters, "test_suffix")
+def test_split_node_id_with_test_suffix(parameters, node_id, expected_class_name, expected_test_name):
+    got = yatest_tools.split_node_id(node_id + parameters, "test_suffix")
+    assert (expected_class_name, expected_test_name + parameters) == got
+
+
+@pytest.mark.parametrize(
+    "node_id,expected_class_name,expected_test_name",
+    [
+        ("/arcadia/data/b/a/test.py::test_b_a", "b.a.test.py", "test_b_a"),
+        ("/arcadia/data/a/test.py::test_a", "a.test.py", "test_a"),
+        ("/arcadia/data/test.py::test", "test.py", "test"),
+    ],
+)
+def test_path_resolving_for_local_conftest_load_policy(
+    mocker, parameters, node_id, expected_class_name, expected_test_name
+):
+    # Order matters
+    extra_modules = [
+        '__tests__.b.a.test',
+        '__tests__.test',
+        '__tests__.a.test',
+    ]
+    mocker.patch.object(sys, 'extra_modules', extra_modules)
+    got = yatest_tools.split_node_id(node_id + parameters)
     assert (expected_class_name, expected_test_name + parameters) == got
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