Add PY_DOCTESTS macro

ISSUE: DEVTOOLS-5060 ([arc::pullid] 223a281f-5913d58c-f1b24f13-cdd71be9) ref:4f0327f2e6ac1221430ac05e15ec977eedd334b0
noisyboy012 · Mar 26, 2019 · 62e1df4 · 62e1df4
1 parent 820af23
commit 62e1df4
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 4 deletions.
diff --git a/build/plugins/pybuild.py b/build/plugins/pybuild.py
@@ -425,6 +425,11 @@ def ontest_srcs(unit, *args):
         unit.onpy_srcs(["NAMESPACE", "__tests__"] + list(args))
 
 
+def onpy_doctests(unit, *args):
+    if unit.get('PY3TEST_BIN' if is_py3(unit) else 'PYTEST_BIN') != 'no':
+        unit.onresource(['-', 'PY_DOCTEST_PACKAGES="{}"'.format(' '.join(args))])
+
+
 def py_register(unit, func, py3):
     if py3:
         unit.on_py3_register([func])

diff --git a/library/python/pytest/main.py b/library/python/pytest/main.py
@@ -1,5 +1,7 @@
 import sys
 
+import __res
+
 
 def main():
     profile = None
@@ -30,6 +32,19 @@ def main():
         if name.startswith(prefix) and not name.endswith('.conftest')
     ]
 
+    doctest_packages = (__res.find("PY_DOCTEST_PACKAGES") or "").split()
+
+    def is_doctest_module(name):
+        for package in doctest_packages:
+            if name == package or name.startswith(package + "."):
+                return True
+        return False
+
+    doctest_modules = [
+        name for name in sys.extra_modules
+        if is_doctest_module(name)
+    ]
+
     def remove_user_site(paths):
         site_paths = ('site-packages', 'site-python')
 
@@ -48,7 +63,7 @@ def is_site_path(path):
 
     sys.path = remove_user_site(sys.path)
     rc = pytest.main(plugins=[
-        collection.CollectionPlugin(test_modules),
+        collection.CollectionPlugin(test_modules, doctest_modules),
         ya,
         conftests,
     ])

diff --git a/library/python/pytest/plugins/collection.py b/library/python/pytest/plugins/collection.py
@@ -9,7 +9,7 @@
 
 class LoadedModule(_pytest.python.Module):
 
-    def __init__(self, name, session):
+    def __init__(self, name, session, namespace=True):
         self.name = name + ".py"  # pytest requires names to be ended with .py
         self.session = session
         self.config = session.config
@@ -18,12 +18,16 @@ def __init__(self, name, session):
         self.keywords = {}
         self.own_markers = []
 
+        self.namespace = namespace
+
     @property
     def nodeid(self):
         return self.name
 
     def _getobj(self):
-        module_name = "__tests__.{}".format(self.name[:-(len(".py"))])
+        module_name = self.name[:-(len(".py"))]
+        if self.namespace:
+            module_name = "__tests__." + module_name
         __import__(module_name)
         return sys.modules[module_name]
 
@@ -45,8 +49,9 @@ def collect(self):
 
 
 class CollectionPlugin(object):
-    def __init__(self, test_modules):
+    def __init__(self, test_modules, doctest_modules):
         self._test_modules = test_modules
+        self._doctest_modules = doctest_modules
 
     def pytest_sessionstart(self, session):
 
@@ -55,4 +60,7 @@ def collect(*args, **kwargs):
                 yield LoadedModule(test_module, session=session)
                 yield DoctestModule(test_module, session=session)
 
+            for doctest_module in self._doctest_modules:
+                yield DoctestModule(doctest_module, session=session, namespace=False)
+
         session.collect = collect