From 222fc079e34c381e76371f316039aaeac717a023 Mon Sep 17 00:00:00 2001
From: Kai <450507+neko-kai@users.noreply.github.com>
Date: Tue, 7 Nov 2023 13:13:27 +0100
Subject: [PATCH] distage-testkit: Don't redundantly run planner multiple times
 for tests that have the same exact dependencies

---
 .../testkit/runner/impl/TestPlanner.scala     | 28 +++++++++++--------
 1 file changed, 17 insertions(+), 11 deletions(-)

diff --git a/distage/distage-testkit-core/src/main/scala/izumi/distage/testkit/runner/impl/TestPlanner.scala b/distage/distage-testkit-core/src/main/scala/izumi/distage/testkit/runner/impl/TestPlanner.scala
index cb3fc01a30..e3885390ad 100644
--- a/distage/distage-testkit-core/src/main/scala/izumi/distage/testkit/runner/impl/TestPlanner.scala
+++ b/distage/distage-testkit-core/src/main/scala/izumi/distage/testkit/runner/impl/TestPlanner.scala
@@ -278,17 +278,23 @@ class TestPlanner[F[_]: TagK: DefaultModule](
       reducedAppModule = appModule.drop(runtimeKeys)
 
       // produce plan for each test
-      testPlans <- tests.map {
-        distageTest =>
-          val forcedRoots = env.forcedRoots.getActiveKeys(fullActivation)
-          val testRoots = distageTest.test.get.diKeys.toSet ++ forcedRoots
-          for {
-            plan <- if (testRoots.nonEmpty) injector.plan(PlannerInput(reducedAppModule, fullActivation, testRoots)) else Right(Plan.empty)
-            _ <- Right(planChecker.showProxyWarnings(plan))
-          } yield {
-            AlmostPreparedTest(distageTest, reducedAppModule, plan.keys, fullActivation)
-          }
-      }.biSequence
+      testPlans <- tests
+        .groupBy {
+          distageTest =>
+            val forcedRoots = env.forcedRoots.getActiveKeys(fullActivation)
+            val testRoots = forcedRoots ++ distageTest.test.get.diKeys
+            testRoots
+        }
+        .toSeq
+        .map {
+          case (testRoots, distageTests) =>
+            for {
+              plan <- if (testRoots.nonEmpty) injector.plan(PlannerInput(reducedAppModule, fullActivation, testRoots)) else Right(Plan.empty)
+              _ <- Right(planChecker.showProxyWarnings(plan))
+            } yield {
+              distageTests.map(AlmostPreparedTest(_, reducedAppModule, plan.keys, fullActivation))
+            }
+        }.biFlatten
       envKeys = testPlans.flatMap(_.targetKeys).toSet
 
       // we need to "strengthen" all _memoized_ weak set instances that occur in our tests to ensure that they