From dbdfddb2423133edefdd5b8dae93fb25f87ca6f5 Mon Sep 17 00:00:00 2001
From: Nikita Kudrin <nikita.kudrin@jetbrains.com>
Date: Fri, 26 Aug 2022 15:04:51 +0300
Subject: [PATCH] [bucketing test] fair bucketing tests

GitOrigin-RevId: 39ddf62ad3851290879636ff681e8880ff605e74
---
 .../core/src/com/intellij/TestCaseLoader.java | 16 +++--
 .../testFramework/FairBucketingTest.java      | 61 +++++++++++++++++++
 2 files changed, 73 insertions(+), 4 deletions(-)
 create mode 100644 platform/testFramework/testSrc/com/intellij/testFramework/FairBucketingTest.java

diff --git a/platform/testFramework/core/src/com/intellij/TestCaseLoader.java b/platform/testFramework/core/src/com/intellij/TestCaseLoader.java
index cb3ff9877e8a..59eaaa1b165d 100644
--- a/platform/testFramework/core/src/com/intellij/TestCaseLoader.java
+++ b/platform/testFramework/core/src/com/intellij/TestCaseLoader.java
@@ -58,7 +58,9 @@ public class TestCaseLoader {
   private static final AtomicInteger CYCLIC_BUCKET_COUNTER = new AtomicInteger(0);
   private static final HashMap<String, Integer> BUCKETS = new HashMap<>();
 
-  /** Split tests into buckets equally across all the buckets */
+  /**
+   * Split tests into buckets equally across all the buckets
+   */
   private static final boolean IS_FAIR_BUCKETING = "true".equals(System.getProperty(FAIR_BUCKETING_FLAG));
 
   /**
@@ -193,18 +195,24 @@ public class TestCaseLoader {
       return MathUtil.nonNegativeAbs(testIdentifier.hashCode()) % TEST_RUNNERS_COUNT == TEST_RUNNER_INDEX;
     }
 
+    return matchesCurrentBucketFair(testIdentifier, TEST_RUNNERS_COUNT, TEST_RUNNER_INDEX);
+  }
+
+  public synchronized static boolean matchesCurrentBucketFair(@NotNull String testIdentifier,
+                                                              int testRunnerCount,
+                                                              int testRunnerIndex) {
     var value = BUCKETS.get(testIdentifier);
 
     if (value != null) {
-      return value == TEST_RUNNER_INDEX;
+      return value == testRunnerIndex;
     }
     else {
       BUCKETS.put(testIdentifier, CYCLIC_BUCKET_COUNTER.getAndIncrement());
     }
 
-    if (CYCLIC_BUCKET_COUNTER.get() == TEST_RUNNERS_COUNT) CYCLIC_BUCKET_COUNTER.set(0);
+    if (CYCLIC_BUCKET_COUNTER.get() == testRunnerCount) CYCLIC_BUCKET_COUNTER.set(0);
 
-    return BUCKETS.get(testIdentifier) == TEST_RUNNER_INDEX;
+    return BUCKETS.get(testIdentifier) == testRunnerIndex;
   }
 
   /**
diff --git a/platform/testFramework/testSrc/com/intellij/testFramework/FairBucketingTest.java b/platform/testFramework/testSrc/com/intellij/testFramework/FairBucketingTest.java
new file mode 100644
index 000000000000..f360300624e8
--- /dev/null
+++ b/platform/testFramework/testSrc/com/intellij/testFramework/FairBucketingTest.java
@@ -0,0 +1,61 @@
+// Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.testFramework;
+
+import com.intellij.TestCaseLoader;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.List;
+
+public class FairBucketingTest {
+
+  private class Data {
+    public String className;
+    public int expectedBucket;
+
+    public Data(String className, int expectedBucket) {
+      this.className = className;
+      this.expectedBucket = expectedBucket;
+    }
+  }
+
+  private int totalBucketsCount = 5;
+  private List<Data> testItems = Arrays.asList(
+    new Data("com.intellij.integrationTests.smoke.idea.IdeaBuildOnJavaTest", 0),
+    new Data("com.intellij.integrationTests.smoke.idea.IdeaBuildOnKotlinTest", 1),
+    new Data("com.intellij.integrationTests.smoke.idea.IdeaCleanImportMavenProjectTest", 2),
+    new Data("com.intellij.integrationTests.smoke.idea.IdeaCreateAllServicesAndExtensionsTest", 3),
+    new Data("com.intellij.integrationTests.smoke.idea.IdeaGenerationTurbochargedSharedIndexesTest", 4),
+    new Data("com.intellij.integrationTests.smoke.idea.IdeaIndexingJavaProjectTest", 0),
+    new Data("com.intellij.integrationTests.smoke.idea.IdeaInspectionOnJavaProjectTest", 1),
+    new Data("com.intellij.integrationTests.smoke.idea.IdeaInspectionOnKotlinProjectTest", 2),
+    new Data("com.intellij.integrationTests.smoke.idea.IdeaInspectionOnSpringProjectTest", 3),
+    new Data("com.intellij.integrationTests.smoke.idea.IdeaJdkClassesCheckOnRedAfterRestartTest", 4),
+    new Data("com.intellij.integrationTests.smoke.idea.IdeaKotlinGradleOpenFilesTest", 0),
+    new Data("com.intellij.integrationTests.smoke.idea.IdeaOpenGradleProjectTest", 1),
+    new Data("com.intellij.integrationTests.smoke.idea.IdeaOpenMavenProjectTest", 2),
+    new Data("com.intellij.integrationTests.smoke.idea.IdeaReopenKotlinProjectFromCacheTest", 3),
+    new Data("com.intellij.integrationTests.smoke.idea.IdeaSpringCloudOpenFilesTest", 4)
+  );
+
+  @Test
+  public void fairBucketingWorks() {
+    for (int currentBucket = 0; currentBucket < totalBucketsCount; currentBucket++) {
+      for (var testData : testItems) {
+        var isMatchedBucket = testData.expectedBucket == currentBucket;
+
+        Assert.assertEquals(String.format("Class `%s` should be in bucket %s", testData.className, currentBucket),
+                            isMatchedBucket,
+                            TestCaseLoader.matchesCurrentBucketFair(testData.className, totalBucketsCount, currentBucket));
+      }
+    }
+  }
+
+  @Test
+  public void multipleFairBucketingInvocation() {
+    fairBucketingWorks();
+    fairBucketingWorks();
+    fairBucketingWorks();
+  }
+}