Fix unit test runner to handle python3 relative imports (PY-17535)

(cherry picked from commit 8728f0fa138b95e6bc1892a74eb8bd595d6cb90a)

python/helpers/pycharm/utrunner.py

@@ -3,6 +3,8 @@ import imp
 import os
 import fnmatch
 
+roots = sys.path[:]
+
 helpers_dir = os.getenv("PYCHARM_HELPERS_DIR", sys.path[0])
 if sys.path[0] != helpers_dir:
   sys.path.insert(0, helpers_dir)
@@ -24,23 +26,50 @@ def loadSource(fileName):
   baseName = os.path.basename(fileName)
   moduleName = os.path.splitext(baseName)[0]
 
+  if os.path.isdir(fileName):
+    fileName = fileName.rstrip('/\\') + os.path.sep
+
   # for users wanted to run unittests under django
-  #because of django took advantage of module name
+  # because of django took advantage of module name
   settings_file = os.getenv('DJANGO_SETTINGS_MODULE')
+
   if settings_file and moduleName == "models":
     baseName = os.path.realpath(fileName)
     moduleName = ".".join((baseName.split(os.sep)[-2], "models"))
+  else:
+    path = fileName
+    for p in roots:
+      # Python 2.6+
+      try:
+        rel_path = os.path.relpath(fileName, start=p)
+        if rel_path.find('..') == -1 and len(rel_path) < len(path):
+          path = rel_path
+      except:
+        pass # relpath can raise an error in case of different drives for a path and start on Windows
 
-  if moduleName in modules and len(sys.argv[1:-1]) == 1: # add unique number to prevent name collisions
+    if path.endswith('.py'):
+      path = path[0:-3]
+
+    moduleName = path.replace('/', '.').replace('\\', '.')
+
+  if moduleName in modules and len(sys.argv[1:-1]) == 1:  # add unique number to prevent name collisions
     cnt = 2
     prefix = moduleName
     while getModuleName(prefix, cnt) in modules:
       cnt += 1
     moduleName = getModuleName(prefix, cnt)
+
   debug("/ Loading " + fileName + " as " + moduleName)
-  if os.path.isdir(fileName):
-    fileName = fileName + os.path.sep
-  module = imp.load_source(moduleName, fileName)
+
+  try:
+    module = imp.load_source(moduleName, fileName)
+  except SystemError:  # probably failed because of the relative imports
+    # first we import module with all its parents
+    __import__(moduleName)
+
+    # then load it by filename to be sure it is the one we need
+    module = imp.load_source(moduleName, fileName)
+
   modules[moduleName] = module
   return module
 
@@ -113,7 +142,7 @@ if __name__ == "__main__":
           debug("/ from folder " + a_splitted[0] + ". Use pattern: " + a_splitted[1])
           modules = loadModulesFromFolderRec(a_splitted[0], a_splitted[1])
       else:
-        if  os.path.isdir(a[0]):
+        if os.path.isdir(a[0]):
           debug("/ from folder " + a[0])
           modules = loadModulesFromFolderRec(a[0])
         else:

python/testData/testRunner/env/unit/relativeImports/relative_imports/tests/test_imps.py

@@ -0,0 +1,7 @@
+import unittest
+
+from ..utils import util
+
+class MyTest(unittest.TestCase):
+  def test_multiply(self):
+    self.assertEquals(4, util.multiply(2, 2))

python/testData/testRunner/env/unit/relativeImports/relative_imports/utils/util.py

@@ -0,0 +1,2 @@
+def multiply(x, y):
+    return x * y

python/testSrc/com/jetbrains/env/PyProcessWithConsoleTestTask.java

@@ -153,7 +153,11 @@ public abstract class PyProcessWithConsoleTestTask<T extends ProcessWithConsoleR
     if (failed.get()) {
       Assert.fail("Failed to run test, see logs for exceptions");
     } else {
-      checkTestResults(runner, stdOut.toString(), stdErr.toString(), stdAll.toString());
+      try {
+        checkTestResults(runner, stdOut.toString(), stdErr.toString(), stdAll.toString());
+      } catch (Throwable e) {
+        throw new RuntimeException(stdAll.toString(), e);
+      }
     }
   }
 

python/testSrc/com/jetbrains/env/python/testing/PythonUnitTestingTest.java

@@ -2,6 +2,7 @@ package com.jetbrains.env.python.testing;
 
 import com.intellij.openapi.util.text.StringUtil;
 import com.jetbrains.env.PyEnvTestCase;
+import com.jetbrains.env.Staging;
 import com.jetbrains.env.ut.PyUnitTestProcessRunner;
 import com.jetbrains.python.psi.LanguageLevel;
 import com.jetbrains.python.testing.PythonTestConfigurationsModel;
@@ -211,4 +212,20 @@ public final class PythonUnitTestingTest extends PyEnvTestCase {
       }
     });
   }
+
+  @Test
+  @Staging
+  public void testRelativeImports() {
+    runPythonTest(new PyUnitTestProcessWithConsoleTestTask("/testRunner/env/unit/relativeImports", "relative_imports/tests/test_imps.py") {
+      @Override
+      protected void checkTestResults(@NotNull final PyUnitTestProcessRunner runner,
+                                      @NotNull final String stdout,
+                                      @NotNull final String stderr,
+                                      @NotNull final String all) {
+        assertEquals(1, runner.getAllTestsCount());
+        assertEquals(1, runner.getPassedTestsCount());
+      }
+    });
+  }
+  
 }