PY-22406: Report unittest as subtests

* Some issues are still missed (see my comments in issue)
* EnvTagsRequired now can be applied to method

python/helpers/pycharm/_jb_runner_tools.py

@@ -205,14 +205,23 @@ class NewTeamcityServiceMessages(_old_service_messages):
             return test_name
 
     # Blocks are used for 2 cases now:
-    # 1) Unittest subtests (broken, because failure can't be reported)
+    # 1) Unittest subtests
     # 2) setup/teardown (does not work, see https://github.com/JetBrains/teamcity-messages/issues/114)
-    # So, temporary disabled
-    # def blockOpened(self, name, flowId=None):
-    #     self.testStarted(".".join(TREE_MANAGER.current_branch + [self._fix_setup_teardown_name(name)]))
-    #
-    # def blockClosed(self, name, flowId=None):
-    #     self.testFinished(".".join(TREE_MANAGER.current_branch + [self._fix_setup_teardown_name(name)]))
+    def blockOpened(self, name, flowId=None):
+         self.testStarted(".".join(TREE_MANAGER.current_branch + [self._fix_setup_teardown_name(name)]))
+
+    def blockClosed(self, name, flowId=None):
+        test_name = ".".join(TREE_MANAGER.current_branch + [self._fix_setup_teardown_name(name)])
+        if self._latest_subtest_result:
+            self.testFinished(test_name)
+        else:
+            self.testFailed(test_name)
+
+        self._latest_subtest_result = None
+
+    def subTestBlockOpened(self, name, subTestResult, flowId=None):
+        self.testStarted(".".join(TREE_MANAGER.current_branch + [self._fix_setup_teardown_name(name)]))
+        self._latest_subtest_result = subTestResult
 
     def testStarted(self, testName, captureStandardOutput=None, flowId=None, is_suite=False):
         test_name_as_list = self._test_to_list(testName)

python/helpers/pycharm/teamcity/messages.py

@@ -81,6 +81,10 @@ class TeamcityServiceMessages(object):
     def blockClosed(self, name, flowId=None):
         self.message('blockClosed', name=name, flowId=flowId)
 
+    # Special PyCharm-specific extension to track subtests, additional property is ignored by TeamCity
+    def subTestBlockOpened(self, name, subTestResult, flowId=None):
+        self.message('blockOpened', name=name, subTestResult=subTestResult, flowId=flowId)
+
     def block(self, name, flowId=None):
         import teamcity.context_managers as cm
         return cm.block(self, name=name, flowId=flowId)

python/helpers/pycharm/teamcity/unittestpy.py

@@ -13,6 +13,7 @@ _real_stdout = sys.stdout
 class TeamcityTestResult(TestResult):
     separator2 = "\n"
 
+    # noinspection PyUnusedLocal
     def __init__(self, stream=_real_stdout, descriptions=None, verbosity=None):
         super(TeamcityTestResult, self).__init__()
 
@@ -111,15 +112,15 @@ class TeamcityTestResult(TestResult):
             self.add_subtest_failure(test_id, block_id)
 
             if issubclass(err[0], test.failureException):
-                self.messages.blockOpened(block_id, flowId=test_id)
+                self.messages.subTestBlockOpened(block_id, subTestResult="Failure", flowId=test_id)
                 self.messages.testStdErr(test_id, out="SubTest failure: %s\n" % convert_error_to_string(err), flowId=test_id)
                 self.messages.blockClosed(block_id, flowId=test_id)
             else:
-                self.messages.blockOpened(block_id, flowId=test_id)
+                self.messages.subTestBlockOpened(block_id, subTestResult="Error", flowId=test_id)
                 self.messages.testStdErr(test_id, out="SubTest error: %s\n" % convert_error_to_string(err), flowId=test_id)
                 self.messages.blockClosed(block_id, flowId=test_id)
         else:
-            self.messages.blockOpened(block_id, flowId=test_id)
+            self.messages.subTestBlockOpened(block_id, subTestResult="Success", flowId=test_id)
             self.messages.blockClosed(block_id, flowId=test_id)
 
     def add_subtest_failure(self, test_id, subtest_block_id):
@@ -195,6 +196,16 @@ class TeamcityTestRunner(TextTestRunner):
         def _makeResult(self):
             return TeamcityTestResult(self.stream, self.descriptions, self.verbosity)
 
+    def run(self, test):
+        # noinspection PyBroadException
+        try:
+            total_tests = test.countTestCases()
+            TeamcityServiceMessages(_real_stdout).testCount(total_tests)
+        except:
+            pass
+
+        return super(TeamcityTestRunner, self).run(test)
+
 
 if __name__ == '__main__':
     from unittest import main

python/testData/testRunner/env/unit/test_subtest.py

@@ -0,0 +1,9 @@
+from unittest import TestCase
+
+
+class SpamTest(TestCase):
+    def test_test(self):
+        for i in range(0, 10):
+            with self.subTest(i=i):
+                print(1)
+                self.assertTrue(i % 2)

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

@@ -30,6 +30,7 @@ import org.junit.runner.Description;
 
 import java.io.File;
 import java.io.IOException;
+import java.lang.reflect.Method;
 import java.util.*;
 
 /**
@@ -236,16 +237,29 @@ public abstract class PyEnvTestCase {
     if (RUN_LOCAL) {
       PyEnvTaskRunner taskRunner = new PyEnvTaskRunner(roots);
 
-      final EnvTestTagsRequired tagsRequiredAnnotation = getClass().getAnnotation(EnvTestTagsRequired.class);
-      final String[] requiredTags;
-      if (tagsRequiredAnnotation != null) {
-        requiredTags = tagsRequiredAnnotation.tags();
+      final EnvTestTagsRequired classAnnotation = getClass().getAnnotation(EnvTestTagsRequired.class);
+      EnvTestTagsRequired methodAnnotation = null;
+      try {
+        final Method method = getClass().getMethod(myTestName.getMethodName());
+        methodAnnotation = method.getAnnotation(EnvTestTagsRequired.class);
       }
-      else {
-        requiredTags = ArrayUtil.EMPTY_STRING_ARRAY;
+      catch (final NoSuchMethodException e) {
+        throw new AssertionError("No such method", e);
       }
+      final String[] classTags = getTags(classAnnotation);
+      final String[] methodTags = getTags(methodAnnotation);
 
-      taskRunner.runTask(testTask, testName, requiredTags);
+      taskRunner.runTask(testTask, testName, ArrayUtil.mergeArrays(methodTags, classTags));
+    }
+  }
+
+  @NotNull
+  private static String[] getTags(@Nullable final EnvTestTagsRequired tagsRequiredAnnotation) {
+    if (tagsRequiredAnnotation != null) {
+      return tagsRequiredAnnotation.tags();
+    }
+    else {
+      return ArrayUtil.EMPTY_STRING_ARRAY;
     }
   }
 

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

@@ -22,6 +22,7 @@ import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.testFramework.EdtTestUtil;
 import com.intellij.testFramework.fixtures.CodeInsightTestFixture;
+import com.jetbrains.env.EnvTestTagsRequired;
 import com.jetbrains.env.PyEnvTestCase;
 import com.jetbrains.env.Staging;
 import com.jetbrains.env.ut.PyScriptTestProcessRunner;
@@ -47,6 +48,36 @@ import static org.junit.Assert.assertEquals;
  */
 public final class PythonUnitTestingTest extends PyEnvTestCase {
 
+  @EnvTestTagsRequired(tags = "python3") // No subtest in py2
+  @Test
+  public void testSubtest() throws Exception {
+    runPythonTest(new PyUnitTestProcessWithConsoleTestTask("testRunner/env/unit/", "test_subtest.py") {
+      @Override
+      protected void checkTestResults(@NotNull PyUnitTestProcessRunner runner,
+                                      @NotNull String stdout,
+                                      @NotNull String stderr,
+                                      @NotNull String all) {
+        final String expectedResult = "Test tree:\n" +
+                                      "[root]\n" +
+                                      ".test_subtest\n" +
+                                      "..SpamTest\n" +
+                                      "...test_test\n" +
+                                      "....(i=0)(-)\n" +
+                                      "....(i=1)(+)\n" +
+                                      "....(i=2)(-)\n" +
+                                      "....(i=3)(+)\n" +
+                                      "....(i=4)(-)\n" +
+                                      "....(i=5)(+)\n" +
+                                      "....(i=6)(-)\n" +
+                                      "....(i=7)(+)\n" +
+                                      "....(i=8)(-)\n" +
+                                      "....(i=9)(+)\n";
+        Assert.assertEquals("", expectedResult, runner.getFormattedTestTree());
+
+      }
+    });
+  }
+
   @Test
   public void testMultipleCases() throws Exception {
     runPythonTest(