From 78e3eea3685229806de98de2b707a86ec7a713e0 Mon Sep 17 00:00:00 2001
From: Alexandr Evstigneev <Alexandr.Evstigneev@jetbrains.com>
Date: Thu, 12 Dec 2019 10:41:18 +0300
Subject: [PATCH] Added bucketing support for cucumber tests

GitOrigin-RevId: 17341382adc6f14ab5fc1dee218aca87b941a2a5
---
 .../com/intellij/cucumber/CucumberMain.java   | 20 ++++++++++++++++-
 .../core/src/com/intellij/TestCaseLoader.java | 22 ++++++++++++++++---
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/build/cucumber-test-runner/src/com/intellij/cucumber/CucumberMain.java b/build/cucumber-test-runner/src/com/intellij/cucumber/CucumberMain.java
index df59d9afe8e4..a7ba3fcdc95d 100644
--- a/build/cucumber-test-runner/src/com/intellij/cucumber/CucumberMain.java
+++ b/build/cucumber-test-runner/src/com/intellij/cucumber/CucumberMain.java
@@ -15,6 +15,7 @@
  */
 package com.intellij.cucumber;
 
+import com.intellij.TestCaseLoader;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.util.Ref;
 import com.intellij.testFramework.TestRunnerUtil;
@@ -22,11 +23,13 @@ import com.intellij.util.ui.UIUtil;
 import cucumber.runtime.Runtime;
 import cucumber.runtime.RuntimeOptions;
 import cucumber.runtime.io.MultiLoader;
+import cucumber.runtime.io.Resource;
 import cucumber.runtime.io.ResourceLoaderClassFinder;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.List;
 
 /**
  * @author Dennis.Ushakov
@@ -57,7 +60,22 @@ public class CucumberMain {
       UIUtil.invokeAndWaitIfNeeded((Runnable)() -> {
         try {
           RuntimeOptions runtimeOptions = new RuntimeOptions(new ArrayList(Arrays.asList(argv)));
-          MultiLoader resourceLoader = new MultiLoader(classLoader);
+          MultiLoader resourceLoader = new MultiLoader(classLoader) {
+            @Override
+            public Iterable<Resource> resources(String path, String suffix) {
+              Iterable<Resource> resources = super.resources(path, suffix);
+              if (TestCaseLoader.shouldBucketTests() && ".feature".equals(suffix)) {
+                List<Resource> filteredResource = new ArrayList<>();
+                resources.forEach(it -> {
+                  if (TestCaseLoader.matchesCurrentBucket(it.getPath())) {
+                    filteredResource.add(it);
+                  }
+                });
+                return filteredResource;
+              }
+              return resources;
+            }
+          };
           ResourceLoaderClassFinder classFinder = new ResourceLoaderClassFinder(resourceLoader, classLoader);
           Runtime runtime = new Runtime(resourceLoader, classFinder, classLoader, runtimeOptions);
           runtimeRef.set(runtime);
diff --git a/platform/testFramework/core/src/com/intellij/TestCaseLoader.java b/platform/testFramework/core/src/com/intellij/TestCaseLoader.java
index 81556ac4a535..24685ed5b266 100644
--- a/platform/testFramework/core/src/com/intellij/TestCaseLoader.java
+++ b/platform/testFramework/core/src/com/intellij/TestCaseLoader.java
@@ -155,14 +155,30 @@ public class TestCaseLoader {
         testCaseClass != myFirstTestClass && testCaseClass != myLastTestClass &&
         TestFrameworkUtil.canRunTest(testCaseClass)) {
 
-      int index = MathUtil.nonNegativeAbs(testCaseClass.getName().hashCode());
-
-      if (index % TEST_RUNNERS_COUNT == TEST_RUNNER_INDEX) {
+      if (matchesCurrentBucket(testCaseClass.getName())) {
         myClassList.add(testCaseClass);
       }
     }
   }
 
+  /**
+   * @return true iff this {@code testIdentifier} matches current testing settings: number of buckets and bucket index. {@code testIdentifier} may
+   * be something identifying a test: test class or feature file name
+   * @apiNote logic for bucketing tests into different bucket configurations.
+   * @see TestCaseLoader#TEST_RUNNERS_COUNT
+   * @see TestCaseLoader#TEST_RUNNER_INDEX
+   */
+  public static boolean matchesCurrentBucket(@NotNull String testIdentifier) {
+    return MathUtil.nonNegativeAbs(testIdentifier.hashCode()) % TEST_RUNNERS_COUNT == TEST_RUNNER_INDEX;
+  }
+
+  /**
+   * @return true iff tests supposed to be separated into buckets using {@link #matchesCurrentBucket(String)} method
+   */
+  public static boolean shouldBucketTests() {
+    return TEST_RUNNERS_COUNT > 1;
+  }
+
   void addFirstTest(Class<?> aClass) {
     assert myFirstTestClass == null : "already added: " + aClass;
     assert shouldAddTestCase(aClass, null, false) : "not a test: " + aClass;