Delegate responsibility/timing of cleaning up old components to caller

2 files changed, 7 insertions, 3 deletions

diff --git a/container-core/src/test/java/com/yahoo/container/di/ContainerTest.java b/container-core/src/test/java/com/yahoo/container/di/ContainerTest.java
index 4ea1321cd70..020aa1e1ae4 100644
--- a/container-core/src/test/java/com/yahoo/container/di/ContainerTest.java
+++ b/container-core/src/test/java/com/yahoo/container/di/ContainerTest.java
@@ -333,11 +333,13 @@ public class ContainerTest extends ContainerTestBase {
     }
 
     ComponentGraph getNewComponentGraph(Container container, ComponentGraph oldGraph) {
-        return container.waitForNextComponentGeneration(oldGraph, Guice.createInjector(), true);
+        Container.ComponentGraphResult result = container.waitForNextComponentGeneration(oldGraph, Guice.createInjector(), true);
+        result.oldComponentsCleanupTask().run();
+        return result.newGraph();
     }
 
     ComponentGraph getNewComponentGraph(Container container) {
-        return container.waitForNextComponentGeneration(new ComponentGraph(), Guice.createInjector(), true);
+        return container.waitForNextComponentGeneration(new ComponentGraph(), Guice.createInjector(), true).newGraph();
     }
 
     private ComponentTakingConfig createComponentTakingConfig(ComponentGraph componentGraph) {
diff --git a/container-core/src/test/java/com/yahoo/container/di/ContainerTestBase.java b/container-core/src/test/java/com/yahoo/container/di/ContainerTestBase.java
index 9565f26fb1e..4f4ffe3732a 100644
--- a/container-core/src/test/java/com/yahoo/container/di/ContainerTestBase.java
+++ b/container-core/src/test/java/com/yahoo/container/di/ContainerTestBase.java
@@ -65,7 +65,9 @@ public class ContainerTestBase {
                             throw new UnsupportedOperationException("getBundle not supported.");
                         }
                     });
-            componentGraph = container.waitForNextComponentGeneration(componentGraph, Guice.createInjector(), true);
+            Container.ComponentGraphResult result = container.waitForNextComponentGeneration(this.componentGraph, Guice.createInjector(), true);
+            result.oldComponentsCleanupTask().run();
+            this.componentGraph = result.newGraph();
         } catch (Exception e) {
             throw new RuntimeException(e);
         }