Schedule NodeAgents using the NodeSpecs from list

5 files changed, 167 insertions, 89 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminImpl.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminImpl.java
index b191401b8e0..801376c23c7 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminImpl.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminImpl.java
@@ -9,6 +9,11 @@ import com.yahoo.vespa.hosted.dockerapi.metrics.MetricReceiverWrapper;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeSpec;
 import com.yahoo.vespa.hosted.node.admin.maintenance.acl.AclMaintainer;
 import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgent;
+import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContext;
+import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContextFactory;
+import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContextManager;
+import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentFactory;
+import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentScheduler;
 import com.yahoo.vespa.hosted.node.admin.util.PrefixLogger;
 
 import java.time.Clock;
@@ -23,7 +28,6 @@ import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 
 /**
@@ -38,7 +42,8 @@ public class NodeAdminImpl implements NodeAdmin {
     private final ScheduledExecutorService metricsScheduler =
             Executors.newScheduledThreadPool(1, ThreadFactoryFactory.getDaemonThreadFactory("metricsscheduler"));
 
-    private final Function<String, NodeAgent> nodeAgentFactory;
+    private final NodeAgentWithSchedulerFactory nodeAgentWithSchedulerFactory;
+    private final NodeAgentContextFactory nodeAgentContextFactory;
     private final Optional<AclMaintainer> aclMaintainer;
 
     private final Clock clock;
@@ -46,16 +51,27 @@ public class NodeAdminImpl implements NodeAdmin {
     private boolean isFrozen;
     private Instant startOfFreezeConvergence;
 
-    private final Map<String, NodeAgent> nodeAgentsByHostname = new ConcurrentHashMap<>();
+    private final Map<String, NodeAgentWithScheduler> nodeAgentWithSchedulerByHostname = new ConcurrentHashMap<>();
 
     private final GaugeWrapper numberOfContainersInLoadImageState;
     private final CounterWrapper numberOfUnhandledExceptionsInNodeAgent;
 
-    public NodeAdminImpl(Function<String, NodeAgent> nodeAgentFactory,
+    public NodeAdminImpl(NodeAgentFactory nodeAgentFactory,
+                         NodeAgentContextFactory nodeAgentContextFactory,
                          Optional<AclMaintainer> aclMaintainer,
                          MetricReceiverWrapper metricReceiver,
                          Clock clock) {
-        this.nodeAgentFactory = nodeAgentFactory;
+        this((NodeAgentWithSchedulerFactory) nodeAgentContext -> create(nodeAgentFactory, nodeAgentContext),
+                nodeAgentContextFactory, aclMaintainer, metricReceiver, clock);
+    }
+
+    NodeAdminImpl(NodeAgentWithSchedulerFactory nodeAgentWithSchedulerFactory,
+                  NodeAgentContextFactory nodeAgentContextFactory,
+                  Optional<AclMaintainer> aclMaintainer,
+                  MetricReceiverWrapper metricReceiver,
+                  Clock clock) {
+        this.nodeAgentWithSchedulerFactory = nodeAgentWithSchedulerFactory;
+        this.nodeAgentContextFactory = nodeAgentContextFactory;
         this.aclMaintainer = aclMaintainer;
 
         this.clock = clock;
@@ -70,22 +86,33 @@ public class NodeAdminImpl implements NodeAdmin {
 
     @Override
     public void refreshContainersToRun(List<NodeSpec> containersToRun) {
-        final Set<String> hostnamesOfContainersToRun = containersToRun.stream()
-                .map(NodeSpec::getHostname)
-                .collect(Collectors.toSet());
+        final Map<String, NodeAgentContext> nodeAgentContextsByHostname = containersToRun.stream()
+                .collect(Collectors.toMap(NodeSpec::getHostname, nodeAgentContextFactory::create));
+
+        // Stop and remove NodeAgents that should no longer be running
+        diff(nodeAgentWithSchedulerByHostname.keySet(), nodeAgentContextsByHostname.keySet())
+                .forEach(hostname -> nodeAgentWithSchedulerByHostname.remove(hostname).stop());
 
-        synchronizeNodesToNodeAgents(hostnamesOfContainersToRun);
+        // Start NodeAgent for hostnames that should be running, but aren't yet
+        diff(nodeAgentContextsByHostname.keySet(), nodeAgentWithSchedulerByHostname.keySet()).forEach(hostname ->  {
+            NodeAgentWithScheduler naws = nodeAgentWithSchedulerFactory.create(nodeAgentContextsByHostname.get(hostname));
+            naws.start();
+            nodeAgentWithSchedulerByHostname.put(hostname, naws);
+        });
 
-        updateNodeAgentMetrics();
+        // At this point, nodeAgentContextsByHostname and nodeAgentWithSchedulerByHostname should have the same keys
+        nodeAgentContextsByHostname.forEach((hostname, context) ->
+            nodeAgentWithSchedulerByHostname.get(hostname).scheduleTickWith(context)
+        );
     }
 
     private void updateNodeAgentMetrics() {
         int numberContainersWaitingImage = 0;
         int numberOfNewUnhandledExceptions = 0;
 
-        for (NodeAgent nodeAgent : nodeAgentsByHostname.values()) {
-            if (nodeAgent.isDownloadingImage()) numberContainersWaitingImage++;
-            numberOfNewUnhandledExceptions += nodeAgent.getAndResetNumberOfUnhandledExceptions();
+        for (NodeAgentWithScheduler nodeAgentWithScheduler : nodeAgentWithSchedulerByHostname.values()) {
+            if (nodeAgentWithScheduler.isDownloadingImage()) numberContainersWaitingImage++;
+            numberOfNewUnhandledExceptions += nodeAgentWithScheduler.getAndResetNumberOfUnhandledExceptions();
         }
 
         numberOfContainersInLoadImageState.sample(numberContainersWaitingImage);
@@ -105,8 +132,8 @@ public class NodeAdminImpl implements NodeAdmin {
         }
 
         // Use filter with count instead of allMatch() because allMatch() will short circuit on first non-match
-        boolean allNodeAgentsConverged = nodeAgentsByHostname.values().stream()
-                .filter(nodeAgent -> !nodeAgent.setFrozen(wantFrozen))
+        boolean allNodeAgentsConverged = nodeAgentWithSchedulerByHostname.values().stream()
+                .filter(nodeAgentScheduler -> !nodeAgentScheduler.setFrozen(wantFrozen))
                 .count() == 0;
 
         if (wantFrozen) {
@@ -134,8 +161,8 @@ public class NodeAdminImpl implements NodeAdmin {
     public void stopNodeAgentServices(List<String> hostnames) {
         // Each container may spend 1-1:30 minutes stopping
         hostnames.parallelStream()
-                .filter(nodeAgentsByHostname::containsKey)
-                .map(nodeAgentsByHostname::get)
+                .filter(nodeAgentWithSchedulerByHostname::containsKey)
+                .map(nodeAgentWithSchedulerByHostname::get)
                 .forEach(nodeAgent -> {
                     nodeAgent.suspend();
                     nodeAgent.stopServices();
@@ -146,7 +173,8 @@ public class NodeAdminImpl implements NodeAdmin {
     public void start() {
         metricsScheduler.scheduleAtFixedRate(() -> {
             try {
-                nodeAgentsByHostname.values().forEach(NodeAgent::updateContainerNodeMetrics);
+                updateNodeAgentMetrics();
+                nodeAgentWithSchedulerByHostname.values().forEach(NodeAgent::updateContainerNodeMetrics);
             } catch (Throwable e) {
                 logger.warning("Metric fetcher scheduler failed", e);
             }
@@ -166,7 +194,7 @@ public class NodeAdminImpl implements NodeAdmin {
         aclScheduler.shutdown();
 
         // Stop all node-agents in parallel, will block until the last NodeAgent is stopped
-        nodeAgentsByHostname.values().parallelStream().forEach(NodeAgent::stop);
+        nodeAgentWithSchedulerByHostname.values().parallelStream().forEach(NodeAgent::stop);
 
         do {
             try {
@@ -185,23 +213,35 @@ public class NodeAdminImpl implements NodeAdmin {
         return result;
     }
 
-    void synchronizeNodesToNodeAgents(Set<String> hostnamesToRun) {
-        // Stop and remove NodeAgents that should no longer be running
-        diff(nodeAgentsByHostname.keySet(), hostnamesToRun)
-                .forEach(hostname -> nodeAgentsByHostname.remove(hostname).stop());
+    static class NodeAgentWithScheduler implements NodeAgent, NodeAgentScheduler {
+        private final NodeAgent nodeAgent;
+        private final NodeAgentScheduler nodeAgentScheduler;
 
-        // Start NodeAgent for hostnames that should be running, but aren't yet
-        diff(hostnamesToRun, nodeAgentsByHostname.keySet())
-                .forEach(this::startNodeAgent);
+        private NodeAgentWithScheduler(NodeAgent nodeAgent, NodeAgentScheduler nodeAgentScheduler) {
+            this.nodeAgent = nodeAgent;
+            this.nodeAgentScheduler = nodeAgentScheduler;
+        }
+
+        @Override public void stopServices() { nodeAgent.stopServices(); }
+        @Override public void suspend() { nodeAgent.suspend(); }
+        @Override public void start() { nodeAgent.start(); }
+        @Override public void stop() { nodeAgent.stop(); }
+        @Override public void updateContainerNodeMetrics() { nodeAgent.updateContainerNodeMetrics(); }
+        @Override public boolean isDownloadingImage() { return nodeAgent.isDownloadingImage(); }
+        @Override public int getAndResetNumberOfUnhandledExceptions() { return nodeAgent.getAndResetNumberOfUnhandledExceptions(); }
+
+        @Override public void scheduleTickWith(NodeAgentContext context) { nodeAgentScheduler.scheduleTickWith(context); }
+        @Override public boolean setFrozen(boolean frozen) { return nodeAgentScheduler.setFrozen(frozen); }
     }
 
-    private void startNodeAgent(String hostname) {
-        if (nodeAgentsByHostname.containsKey(hostname))
-            throw new IllegalArgumentException("Attempted to start NodeAgent for hostname " + hostname +
-                    ", but one is already running!");
+    @FunctionalInterface
+    interface NodeAgentWithSchedulerFactory {
+        NodeAgentWithScheduler create(NodeAgentContext context);
+    }
 
-        NodeAgent agent = nodeAgentFactory.apply(hostname);
-        agent.start();
-        nodeAgentsByHostname.put(hostname, agent);
+    private static NodeAgentWithScheduler create(NodeAgentFactory nodeAgentFactory, NodeAgentContext context) {
+        NodeAgentContextManager contextManager = new NodeAgentContextManager(context);
+        NodeAgent nodeAgent = nodeAgentFactory.create(contextManager);
+        return new NodeAgentWithScheduler(nodeAgent, contextManager);
     }
 }
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextFactory.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextFactory.java
new file mode 100644
index 00000000000..0cfafe34717
--- /dev/null
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextFactory.java
@@ -0,0 +1,12 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.node.admin.nodeagent;
+
+import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeSpec;
+
+/**
+ * @author freva
+ */
+@FunctionalInterface
+public interface NodeAgentContextFactory {
+    NodeAgentContext create(NodeSpec nodeSpec);
+}
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentFactory.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentFactory.java
new file mode 100644
index 00000000000..bd13b7eb094
--- /dev/null
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentFactory.java
@@ -0,0 +1,10 @@
+// Copyright 2019 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.node.admin.nodeagent;
+
+/**
+ * @author freva
+ */
+@FunctionalInterface
+public interface NodeAgentFactory {
+    NodeAgent create(NodeAgentContextSupplier contextSupplier);
+}
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/DockerTester.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/DockerTester.java
index cb3a1fb5e2c..109bce4c13f 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/DockerTester.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/DockerTester.java
@@ -15,8 +15,9 @@ import com.yahoo.vespa.hosted.node.admin.docker.DockerOperationsImpl;
 import com.yahoo.vespa.hosted.node.admin.maintenance.StorageMaintainer;
 import com.yahoo.vespa.hosted.node.admin.nodeadmin.NodeAdminImpl;
 import com.yahoo.vespa.hosted.node.admin.nodeadmin.NodeAdminStateUpdater;
-import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgent;
+import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContextFactory;
 import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContextImpl;
+import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentFactory;
 import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentImpl;
 import com.yahoo.vespa.hosted.node.admin.task.util.network.IPAddressesMock;
 import com.yahoo.vespa.hosted.provision.Node;
@@ -30,7 +31,6 @@ import java.nio.file.Paths;
 import java.time.Clock;
 import java.time.Duration;
 import java.util.Optional;
-import java.util.function.Function;
 import java.util.logging.Logger;
 
 import static com.yahoo.yolean.Exceptions.uncheck;
@@ -87,19 +87,20 @@ public class DockerTester implements AutoCloseable {
                 .build();
         nodeRepository.updateNodeRepositoryNode(hostSpec);
 
-        Clock clock = Clock.systemUTC();
         FileSystem fileSystem = TestFileSystem.create();
         DockerOperations dockerOperations = new DockerOperationsImpl(docker, processExecuter, ipAddresses);
 
         MetricReceiverWrapper mr = new MetricReceiverWrapper(MetricReceiver.nullImplementation);
-        Function<String, NodeAgent> nodeAgentFactory = (hostName) -> new NodeAgentImpl(
-                new NodeAgentContextImpl.Builder(hostName).fileSystem(fileSystem).build(), nodeRepository,
-                orchestrator, dockerOperations, storageMaintainer, clock, INTERVAL, Optional.empty(), Optional.empty(), Optional.empty());
-        nodeAdmin = new NodeAdminImpl(nodeAgentFactory, Optional.empty(), mr, Clock.systemUTC());
+        NodeAgentFactory nodeAgentFactory = contextSupplier -> new NodeAgentImpl(
+                contextSupplier, nodeRepository,
+                orchestrator, dockerOperations, storageMaintainer, Optional.empty(), Optional.empty(), Optional.empty());
+        NodeAgentContextFactory nodeAgentContextFactory = nodeSpec ->
+                new NodeAgentContextImpl.Builder(nodeSpec).fileSystem(fileSystem).build();
+        nodeAdmin = new NodeAdminImpl(nodeAgentFactory, nodeAgentContextFactory, Optional.empty(), mr, Clock.systemUTC());
         nodeAdminStateUpdater = new NodeAdminStateUpdater(nodeRepository, orchestrator,
                 nodeAdmin, HOST_HOSTNAME);
 
-        this.loopThread = new Thread(() -> {
+        loopThread = new Thread(() -> {
             nodeAdminStateUpdater.start();
 
             while (! terminated) {
@@ -135,8 +136,10 @@ public class DockerTester implements AutoCloseable {
 
     @Override
     public void close() {
-        terminated = true;
+        // First, stop NodeAdmin and all the NodeAgents
+        nodeAdmin.stop();
 
+        terminated = true;
         do {
             try {
                 loopThread.join();
@@ -144,8 +147,5 @@ public class DockerTester implements AutoCloseable {
                 e.printStackTrace();
             }
         } while (loopThread.isAlive());
-
-        // Finally, stop NodeAdmin and all the NodeAgents
-        nodeAdmin.stop();
     }
 }
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminImplTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminImplTest.java
index 3860e2e9780..766638b94cb 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminImplTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminImplTest.java
@@ -1,22 +1,23 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.node.admin.nodeadmin;
 
+import com.yahoo.config.provision.NodeType;
 import com.yahoo.metrics.simple.MetricReceiver;
 import com.yahoo.test.ManualClock;
 import com.yahoo.vespa.hosted.dockerapi.metrics.MetricReceiverWrapper;
-import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgent;
-import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentImpl;
+import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeSpec;
+import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContext;
+import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContextFactory;
+import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContextImpl;
+import com.yahoo.vespa.hosted.provision.Node;
 import org.junit.Test;
 import org.mockito.InOrder;
 
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
-import java.util.Set;
-import java.util.function.Function;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -31,75 +32,72 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
+import static com.yahoo.vespa.hosted.node.admin.nodeadmin.NodeAdminImpl.NodeAgentWithScheduler;
+import static com.yahoo.vespa.hosted.node.admin.nodeadmin.NodeAdminImpl.NodeAgentWithSchedulerFactory;
+
 /**
  * @author bakksjo
  */
 public class NodeAdminImplTest {
-    // Trick to allow mocking of typed interface without casts/warnings.
-    private interface NodeAgentFactory extends Function<String, NodeAgent> {}
-    private final Function<String, NodeAgent> nodeAgentFactory = mock(NodeAgentFactory.class);
+
+    private final NodeAgentWithSchedulerFactory nodeAgentWithSchedulerFactory = mock(NodeAgentWithSchedulerFactory.class);
+    private final NodeAgentContextFactory nodeAgentContextFactory = mock(NodeAgentContextFactory.class);
     private final ManualClock clock = new ManualClock();
 
-    private final NodeAdminImpl nodeAdmin = new NodeAdminImpl(nodeAgentFactory, Optional.empty(),
-            new MetricReceiverWrapper(MetricReceiver.nullImplementation), clock);
+    private final NodeAdminImpl nodeAdmin = new NodeAdminImpl(nodeAgentWithSchedulerFactory, nodeAgentContextFactory,
+            Optional.empty(), new MetricReceiverWrapper(MetricReceiver.nullImplementation), clock);
 
     @Test
     public void nodeAgentsAreProperlyLifeCycleManaged() {
-        final String hostName1 = "host1.test.yahoo.com";
-        final String hostName2 = "host2.test.yahoo.com";
-        final NodeAgent nodeAgent1 = mock(NodeAgentImpl.class);
-        final NodeAgent nodeAgent2 = mock(NodeAgentImpl.class);
-        when(nodeAgentFactory.apply(eq(hostName1))).thenReturn(nodeAgent1);
-        when(nodeAgentFactory.apply(eq(hostName2))).thenReturn(nodeAgent2);
+        final NodeSpec nodeSpec1 = createNodeSpec("host1.test.yahoo.com");
+        final NodeSpec nodeSpec2 = createNodeSpec("host2.test.yahoo.com");
+        final NodeAgentWithScheduler nodeAgent1 = mockNodeAgentWithSchedulerFactory(nodeSpec1);
+        final NodeAgentWithScheduler nodeAgent2 = mockNodeAgentWithSchedulerFactory(nodeSpec2);
 
+        final InOrder inOrder = inOrder(nodeAgentWithSchedulerFactory, nodeAgent1, nodeAgent2);
+        nodeAdmin.refreshContainersToRun(Collections.emptyList());
+        verifyNoMoreInteractions(nodeAgentWithSchedulerFactory);
 
-        final InOrder inOrder = inOrder(nodeAgentFactory, nodeAgent1, nodeAgent2);
-        nodeAdmin.synchronizeNodesToNodeAgents(Collections.emptySet());
-        verifyNoMoreInteractions(nodeAgentFactory);
-
-        nodeAdmin.synchronizeNodesToNodeAgents(Collections.singleton(hostName1));
-        inOrder.verify(nodeAgentFactory).apply(hostName1);
+        nodeAdmin.refreshContainersToRun(Collections.singletonList(nodeSpec1));
         inOrder.verify(nodeAgent1).start();
+        inOrder.verify(nodeAgent2, never()).start();
         inOrder.verify(nodeAgent1, never()).stop();
 
-        nodeAdmin.synchronizeNodesToNodeAgents(Collections.singleton(hostName1));
-        inOrder.verify(nodeAgentFactory, never()).apply(any(String.class));
+        nodeAdmin.refreshContainersToRun(Collections.singletonList(nodeSpec1));
+        inOrder.verify(nodeAgentWithSchedulerFactory, never()).create(any());
         inOrder.verify(nodeAgent1, never()).start();
         inOrder.verify(nodeAgent1, never()).stop();
 
-        nodeAdmin.synchronizeNodesToNodeAgents(Collections.emptySet());
-        inOrder.verify(nodeAgentFactory, never()).apply(any(String.class));
+        nodeAdmin.refreshContainersToRun(Collections.emptyList());
+        inOrder.verify(nodeAgentWithSchedulerFactory, never()).create(any());
         verify(nodeAgent1).stop();
 
-        nodeAdmin.synchronizeNodesToNodeAgents(Collections.singleton(hostName2));
-        inOrder.verify(nodeAgentFactory).apply(hostName2);
+        nodeAdmin.refreshContainersToRun(Collections.singletonList(nodeSpec2));
         inOrder.verify(nodeAgent2).start();
         inOrder.verify(nodeAgent2, never()).stop();
-        verify(nodeAgent1).stop();
+        inOrder.verify(nodeAgent1, never()).stop();
 
-        nodeAdmin.synchronizeNodesToNodeAgents(Collections.emptySet());
-        inOrder.verify(nodeAgentFactory, never()).apply(any(String.class));
+        nodeAdmin.refreshContainersToRun(Collections.emptyList());
+        inOrder.verify(nodeAgentWithSchedulerFactory, never()).create(any());
         inOrder.verify(nodeAgent2, never()).start();
         inOrder.verify(nodeAgent2).stop();
-
-        verifyNoMoreInteractions(nodeAgent1);
-        verifyNoMoreInteractions(nodeAgent2);
+        inOrder.verify(nodeAgent1, never()).start();
+        inOrder.verify(nodeAgent1, never()).stop();
     }
 
     @Test
     public void testSetFrozen() {
-        List<NodeAgent> nodeAgents = new ArrayList<>();
-        Set<String> existingContainerHostnames = new HashSet<>();
+        List<NodeSpec> nodeSpecs = new ArrayList<>();
+        List<NodeAgentWithScheduler> nodeAgents = new ArrayList<>();
         for (int i = 0; i < 3; i++) {
-            final String hostName = "host" + i + ".test.yahoo.com";
-            NodeAgent nodeAgent = mock(NodeAgent.class);
-            nodeAgents.add(nodeAgent);
-            when(nodeAgentFactory.apply(eq(hostName))).thenReturn(nodeAgent);
+            NodeSpec nodeSpec = createNodeSpec("host" + i + ".test.yahoo.com");
+            NodeAgentWithScheduler nodeAgent = mockNodeAgentWithSchedulerFactory(nodeSpec);
 
-            existingContainerHostnames.add(hostName);
+            nodeSpecs.add(nodeSpec);
+            nodeAgents.add(nodeAgent);
         }
 
-        nodeAdmin.synchronizeNodesToNodeAgents(existingContainerHostnames);
+        nodeAdmin.refreshContainersToRun(nodeSpecs);
 
         assertTrue(nodeAdmin.isFrozen()); // Initially everything is frozen to force convergence
         mockNodeAgentSetFrozenResponse(nodeAgents, true, true, true);
@@ -155,10 +153,28 @@ public class NodeAdminImplTest {
         assertEquals(Duration.ofSeconds(1), nodeAdmin.subsystemFreezeDuration());
     }
 
-    private void mockNodeAgentSetFrozenResponse(List<NodeAgent> nodeAgents, boolean... responses) {
+    private void mockNodeAgentSetFrozenResponse(List<NodeAgentWithScheduler> nodeAgents, boolean... responses) {
         for (int i = 0; i < nodeAgents.size(); i++) {
-            NodeAgent nodeAgent = nodeAgents.get(i);
+            NodeAgentWithScheduler nodeAgent = nodeAgents.get(i);
             when(nodeAgent.setFrozen(anyBoolean())).thenReturn(responses[i]);
         }
     }
+
+    private NodeSpec createNodeSpec(String hostname) {
+        return new NodeSpec.Builder()
+                .hostname(hostname)
+                .state(Node.State.active)
+                .nodeType(NodeType.tenant)
+                .flavor("default")
+                .build();
+    }
+
+    private NodeAgentWithScheduler mockNodeAgentWithSchedulerFactory(NodeSpec nodeSpec) {
+        NodeAgentContext context = new NodeAgentContextImpl.Builder(nodeSpec).build();
+        when(nodeAgentContextFactory.create(eq(nodeSpec))).thenReturn(context);
+
+        NodeAgentWithScheduler nodeAgentWithScheduler = mock(NodeAgentWithScheduler.class);
+        when(nodeAgentWithSchedulerFactory.create(eq(context))).thenReturn(nodeAgentWithScheduler);
+        return nodeAgentWithScheduler;
+    }
 }