Simplify NodeSpec creation for testing

9 files changed, 72 insertions, 161 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeSpec.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeSpec.java
index 0e8caf728f1..09496b88aeb 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeSpec.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeSpec.java
@@ -9,6 +9,7 @@ import com.yahoo.config.provision.NodeResources;
 import com.yahoo.config.provision.NodeType;
 
 import java.time.Instant;
+import java.util.EnumSet;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
@@ -682,5 +683,34 @@ public class NodeSpec {
                     reports, parentHostname);
         }
 
+
+        public static Builder testSpec(String hostname) {
+            return testSpec(hostname, NodeState.active);
+        }
+
+        /**
+         * Creates a NodeSpec.Builder that has the given hostname, in a given state, and some
+         * reasonable values for the remaining required NodeSpec fields.
+         */
+        public static Builder testSpec(String hostname, NodeState state) {
+            Builder builder = new Builder()
+                    .hostname(hostname)
+                    .state(state)
+                    .type(NodeType.tenant)
+                    .flavor("d-2-8-50")
+                    .resources(new NodeResources(2, 8, 50, 10));
+
+            // Set the required allocated fields
+            if (EnumSet.of(NodeState.active, NodeState.inactive, NodeState.reserved).contains(state)) {
+                builder .owner(ApplicationId.defaultId())
+                        .membership(new NodeMembership("container", "my-id", "group", 0, false))
+                        .wantedVespaVersion(Version.fromString("7.1.1"))
+                        .wantedDockerImage(DockerImage.fromString("docker.domain.tld/repo/image:7.1.1"))
+                        .currentRestartGeneration(0)
+                        .wantedRestartGeneration(0);
+            }
+
+            return builder;
+        }
     }
 }
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextImpl.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextImpl.java
index 02e01264a57..4585d72d6de 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextImpl.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextImpl.java
@@ -2,7 +2,6 @@
 package com.yahoo.vespa.hosted.node.admin.nodeagent;
 
 import com.yahoo.config.provision.CloudName;
-import com.yahoo.config.provision.NodeType;
 import com.yahoo.config.provision.SystemName;
 import com.yahoo.config.provision.zone.ZoneApi;
 import com.yahoo.config.provision.zone.ZoneId;
@@ -11,7 +10,6 @@ import com.yahoo.vespa.athenz.api.AthenzService;
 import com.yahoo.vespa.hosted.dockerapi.ContainerName;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.Acl;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeSpec;
-import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeState;
 import com.yahoo.vespa.hosted.node.admin.docker.DockerNetworking;
 
 import java.nio.file.FileSystem;
@@ -187,7 +185,7 @@ public class NodeAgentContextImpl implements NodeAgentContext {
 
     /** For testing only! */
     public static class Builder {
-        private NodeSpec.Builder nodeSpecBuilder = new NodeSpec.Builder();
+        private NodeSpec.Builder nodeSpecBuilder;
         private Acl acl;
         private AthenzIdentity identity;
         private DockerNetworking dockerNetworking;
@@ -209,11 +207,7 @@ public class NodeAgentContextImpl implements NodeAgentContext {
          * if you want to control the entire NodeSpec.
          */
         public Builder(String hostname) {
-            this.nodeSpecBuilder
-                    .hostname(hostname)
-                    .state(NodeState.active)
-                    .type(NodeType.tenant)
-                    .flavor("d-2-8-50");
+            this.nodeSpecBuilder = NodeSpec.Builder.testSpec(hostname);
         }
 
         public Builder nodeSpecBuilder(Function<NodeSpec.Builder, NodeSpec.Builder> nodeSpecBuilderModifier) {
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/DockerFailTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/DockerFailTest.java
index 7776bf1f224..0ed52cae5b8 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/DockerFailTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/DockerFailTest.java
@@ -2,10 +2,8 @@
 package com.yahoo.vespa.hosted.node.admin.integrationTests;
 
 import com.yahoo.config.provision.DockerImage;
-import com.yahoo.config.provision.NodeType;
 import com.yahoo.vespa.hosted.dockerapi.ContainerName;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeSpec;
-import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeState;
 import org.junit.Test;
 
 import static org.mockito.ArgumentMatchers.any;
@@ -24,18 +22,10 @@ public class DockerFailTest {
             final DockerImage dockerImage = DockerImage.fromString("dockerImage");
             final ContainerName containerName = new ContainerName("host1");
             final String hostname = "host1.test.yahoo.com";
-            tester.addChildNodeRepositoryNode(new NodeSpec.Builder()
-                    .hostname(hostname)
+            tester.addChildNodeRepositoryNode(NodeSpec.Builder
+                    .testSpec(hostname)
                     .wantedDockerImage(dockerImage)
                     .currentDockerImage(dockerImage)
-                    .state(NodeState.active)
-                    .type(NodeType.tenant)
-                    .flavor("docker")
-                    .wantedRestartGeneration(1L)
-                    .currentRestartGeneration(1L)
-                    .vcpu(1)
-                    .memoryGb(1)
-                    .diskGb(1)
                     .build());
 
             tester.inOrder(tester.docker).createContainerCommand(eq(dockerImage), eq(containerName));
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 ea41ab21b9d..9b509829ba2 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
@@ -7,7 +7,6 @@ import com.yahoo.vespa.flags.InMemoryFlagSource;
 import com.yahoo.vespa.hosted.dockerapi.Docker;
 import com.yahoo.vespa.hosted.dockerapi.metrics.Metrics;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeSpec;
-import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeState;
 import com.yahoo.vespa.hosted.node.admin.configserver.orchestrator.Orchestrator;
 import com.yahoo.vespa.hosted.node.admin.docker.DockerOperations;
 import com.yahoo.vespa.hosted.node.admin.docker.DockerOperationsImpl;
@@ -76,14 +75,7 @@ public class DockerTester implements AutoCloseable {
 
         TerminalImpl terminal = new TerminalImpl(command -> new TestChildProcess2(0, ""));
 
-        NodeSpec hostSpec = new NodeSpec.Builder()
-                .hostname(HOST_HOSTNAME.value())
-                .state(NodeState.active)
-                .type(NodeType.host)
-                .flavor("default")
-                .wantedRestartGeneration(1L)
-                .currentRestartGeneration(1L)
-                .build();
+        NodeSpec hostSpec = NodeSpec.Builder.testSpec(HOST_HOSTNAME.value()).type(NodeType.host).build();
         nodeRepository.updateNodeRepositoryNode(hostSpec);
 
         Clock clock = Clock.systemUTC();
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/MultiDockerTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/MultiDockerTest.java
index a09c6a9e907..e5156e68810 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/MultiDockerTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/MultiDockerTest.java
@@ -2,7 +2,6 @@
 package com.yahoo.vespa.hosted.node.admin.integrationTests;
 
 import com.yahoo.config.provision.DockerImage;
-import com.yahoo.config.provision.NodeType;
 import com.yahoo.vespa.hosted.dockerapi.ContainerName;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeAttributes;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeSpec;
@@ -25,13 +24,7 @@ public class MultiDockerTest {
             NodeSpec nodeSpec2 = addAndWaitForNode(
                     tester, "host2.test.yahoo.com", DockerImage.fromString("image2"));
 
-            tester.addChildNodeRepositoryNode(
-                    new NodeSpec.Builder(nodeSpec2)
-                            .state(NodeState.dirty)
-                            .vcpu(1)
-                            .memoryGb(1)
-                            .diskGb(1)
-                            .build());
+            tester.addChildNodeRepositoryNode(NodeSpec.Builder.testSpec(nodeSpec2.hostname(), NodeState.dirty).build());
 
             tester.inOrder(tester.docker).deleteContainer(eq(new ContainerName("host2")));
             tester.inOrder(tester.storageMaintainer).archiveNodeStorage(
@@ -43,19 +36,7 @@ public class MultiDockerTest {
     }
 
     private NodeSpec addAndWaitForNode(DockerTester tester, String hostName, DockerImage dockerImage) {
-        NodeSpec nodeSpec = new NodeSpec.Builder()
-                .hostname(hostName)
-                .wantedDockerImage(dockerImage)
-                .state(NodeState.active)
-                .type(NodeType.tenant)
-                .flavor("docker")
-                .wantedRestartGeneration(1L)
-                .currentRestartGeneration(1L)
-                .vcpu(2)
-                .memoryGb(4)
-                .diskGb(1)
-                .build();
-
+        NodeSpec nodeSpec = NodeSpec.Builder.testSpec(hostName).wantedDockerImage(dockerImage).build();
         tester.addChildNodeRepositoryNode(nodeSpec);
 
         ContainerName containerName = ContainerName.fromHostname(hostName);
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/RebootTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/RebootTest.java
index d1778982043..30f23de15bf 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/RebootTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/RebootTest.java
@@ -1,12 +1,9 @@
 // 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.integrationTests;
 
-import com.yahoo.component.Version;
 import com.yahoo.config.provision.DockerImage;
-import com.yahoo.config.provision.NodeType;
 import com.yahoo.vespa.hosted.dockerapi.ContainerName;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeSpec;
-import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeState;
 import com.yahoo.vespa.hosted.node.admin.nodeadmin.NodeAdminStateUpdater;
 import org.junit.Test;
 
@@ -31,7 +28,7 @@ public class RebootTest {
     @Test
     public void test() {
         try (DockerTester tester = new DockerTester()) {
-            tester.addChildNodeRepositoryNode(createNodeRepositoryNode());
+            tester.addChildNodeRepositoryNode(NodeSpec.Builder.testSpec(hostname).wantedDockerImage(dockerImage).build());
 
             tester.inOrder(tester.docker).createContainerCommand(eq(dockerImage), eq(new ContainerName("host1")));
 
@@ -46,17 +43,4 @@ public class RebootTest {
             assertTrue(tester.nodeAdmin.setFrozen(true));
         }
     }
-
-    private NodeSpec createNodeRepositoryNode() {
-        return new NodeSpec.Builder()
-                .hostname(hostname)
-                .wantedDockerImage(dockerImage)
-                .state(NodeState.active)
-                .type(NodeType.tenant)
-                .flavor("docker")
-                .currentVespaVersion(Version.fromString("6.50.0"))
-                .wantedRestartGeneration(1L)
-                .currentRestartGeneration(1L)
-                .build();
-    }
 }
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/RestartTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/RestartTest.java
index bfc54cac045..dd727d102e0 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/RestartTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/RestartTest.java
@@ -2,11 +2,9 @@
 package com.yahoo.vespa.hosted.node.admin.integrationTests;
 
 import com.yahoo.config.provision.DockerImage;
-import com.yahoo.config.provision.NodeType;
 import com.yahoo.vespa.hosted.dockerapi.ContainerName;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeAttributes;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeSpec;
-import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeState;
 import org.junit.Test;
 
 import static com.yahoo.vespa.hosted.node.admin.integrationTests.DockerTester.NODE_PROGRAM;
@@ -26,15 +24,7 @@ public class RestartTest {
             String hostname = "host1.test.yahoo.com";
             DockerImage dockerImage = DockerImage.fromString("dockerImage:1.2.3");
 
-            tester.addChildNodeRepositoryNode(new NodeSpec.Builder()
-                    .hostname(hostname)
-                    .state(NodeState.active)
-                    .wantedDockerImage(dockerImage)
-                    .type(NodeType.tenant)
-                    .flavor("docker")
-                    .wantedRestartGeneration(1)
-                    .currentRestartGeneration(1)
-                    .build());
+            tester.addChildNodeRepositoryNode(NodeSpec.Builder.testSpec(hostname).wantedDockerImage(dockerImage).build());
 
             tester.inOrder(tester.docker).createContainerCommand(eq(dockerImage), eq(new ContainerName("host1")));
             tester.inOrder(tester.nodeRepository).updateNodeAttributes(
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminStateUpdaterTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminStateUpdaterTest.java
index 9a25e4188f9..36530ff014c 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminStateUpdaterTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminStateUpdaterTest.java
@@ -2,7 +2,6 @@
 package com.yahoo.vespa.hosted.node.admin.nodeadmin;
 
 import com.yahoo.config.provision.HostName;
-import com.yahoo.config.provision.NodeType;
 import com.yahoo.test.ManualClock;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.Acl;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeRepository;
@@ -258,28 +257,12 @@ public class NodeAdminStateUpdaterTest {
 
     private void mockNodeRepo(NodeState hostState, int numberOfNodes) {
         List<NodeSpec> containersToRun = IntStream.range(1, numberOfNodes + 1)
-                .mapToObj(i -> new NodeSpec.Builder()
-                        .hostname("host" + i + ".yahoo.com")
-                        .state(NodeState.active)
-                        .type(NodeType.tenant)
-                        .flavor("docker")
-                        .vcpu(1)
-                        .memoryGb(1)
-                        .diskGb(1)
-                        .build())
+                .mapToObj(i -> NodeSpec.Builder.testSpec("host" + i + ".yahoo.com").build())
                 .collect(Collectors.toList());
 
         when(nodeRepository.getNodes(eq(hostHostname.value()))).thenReturn(containersToRun);
-
-        when(nodeRepository.getNode(eq(hostHostname.value()))).thenReturn(new NodeSpec.Builder()
-                .hostname(hostHostname.value())
-                .state(hostState)
-                .type(NodeType.tenant)
-                .flavor("default")
-                .vcpu(1)
-                .memoryGb(1)
-                .diskGb(1)
-                .build());
+        when(nodeRepository.getNode(eq(hostHostname.value()))).thenReturn(
+                NodeSpec.Builder.testSpec(hostHostname.value(), hostState).build());
     }
 
     private void mockAcl(Acl acl, int... nodeIds) {
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImplTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImplTest.java
index 06d15a3f577..b895300d7b1 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImplTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImplTest.java
@@ -52,13 +52,7 @@ import static org.mockito.Mockito.when;
 public class NodeAgentImplTest {
     private static final NodeResources resources = new NodeResources(2, 16, 250, 1, NodeResources.DiskSpeed.fast, NodeResources.StorageType.local);
     private static final Version vespaVersion = Version.fromString("1.2.3");
-
-    private final String hostName = "host1.test.yahoo.com";
-    private final NodeSpec.Builder nodeBuilder = new NodeSpec.Builder()
-            .hostname(hostName)
-            .type(NodeType.tenant)
-            .flavor("docker")
-            .resources(resources);
+    private static final String hostName = "host1.test.yahoo.com";
 
     private final NodeAgentContextSupplier contextSupplier = mock(NodeAgentContextSupplier.class);
     private final DockerImage dockerImage = DockerImage.fromString("dockerImage");
@@ -75,11 +69,9 @@ public class NodeAgentImplTest {
 
     @Test
     public void upToDateContainerIsUntouched() {
-        final NodeSpec node = nodeBuilder
-                .state(NodeState.active)
+        final NodeSpec node = nodeBuilder(NodeState.active)
                 .wantedDockerImage(dockerImage).currentDockerImage(dockerImage)
                 .wantedVespaVersion(vespaVersion).currentVespaVersion(vespaVersion)
-                .wantedRestartGeneration(1).currentRestartGeneration(1)
                 .allowedToBeDown(false)
                 .build();
 
@@ -104,11 +96,9 @@ public class NodeAgentImplTest {
 
     @Test
     public void verifyRemoveOldFilesIfDiskFull() {
-        final NodeSpec node = nodeBuilder
-                .state(NodeState.active)
+        final NodeSpec node = nodeBuilder(NodeState.active)
                 .wantedDockerImage(dockerImage).currentDockerImage(dockerImage)
                 .wantedVespaVersion(vespaVersion).currentVespaVersion(vespaVersion)
-                .wantedRestartGeneration(1).currentRestartGeneration(1)
                 .build();
 
         NodeAgentContext context = createContext(node);
@@ -124,11 +114,9 @@ public class NodeAgentImplTest {
     @Test
     public void startsAfterStoppingServices() {
         final InOrder inOrder = inOrder(dockerOperations);
-        final NodeSpec node = nodeBuilder
-                .state(NodeState.active)
+        final NodeSpec node = nodeBuilder(NodeState.active)
                 .wantedDockerImage(dockerImage).currentDockerImage(dockerImage)
                 .wantedVespaVersion(vespaVersion).currentVespaVersion(vespaVersion)
-                .wantedRestartGeneration(1).currentRestartGeneration(1)
                 .build();
 
         NodeAgentContext context = createContext(node);
@@ -160,11 +148,9 @@ public class NodeAgentImplTest {
 
     @Test
     public void absentContainerCausesStart() {
-        final NodeSpec node = nodeBuilder
-                .state(NodeState.active)
+        final NodeSpec node = nodeBuilder(NodeState.active)
                 .wantedDockerImage(dockerImage)
                 .wantedVespaVersion(vespaVersion)
-                .wantedRestartGeneration(1).currentRestartGeneration(1)
                 .build();
 
         NodeAgentContext context = createContext(node);
@@ -195,11 +181,9 @@ public class NodeAgentImplTest {
     @Test
     public void containerIsNotStoppedIfNewImageMustBePulled() {
         final DockerImage newDockerImage = DockerImage.fromString("new-image");
-        final NodeSpec node = nodeBuilder
-                .state(NodeState.active)
+        final NodeSpec node = nodeBuilder(NodeState.active)
                 .wantedDockerImage(newDockerImage).currentDockerImage(dockerImage)
                 .wantedVespaVersion(vespaVersion).currentVespaVersion(vespaVersion)
-                .wantedRestartGeneration(1).currentRestartGeneration(1)
                 .build();
 
         NodeAgentContext context = createContext(node);
@@ -221,11 +205,9 @@ public class NodeAgentImplTest {
 
     @Test
     public void containerIsUpdatedIfCpuChanged() {
-        NodeSpec.Builder specBuilder = nodeBuilder
-                .state(NodeState.active)
+        NodeSpec.Builder specBuilder = nodeBuilder(NodeState.active)
                 .wantedDockerImage(dockerImage).currentDockerImage(dockerImage)
                 .wantedVespaVersion(vespaVersion).currentVespaVersion(vespaVersion)
-                .wantedRestartGeneration(1).currentRestartGeneration(1)
                 .allowedToBeDown(false);
 
         NodeAgentContext firstContext = createContext(specBuilder.build());
@@ -271,8 +253,7 @@ public class NodeAgentImplTest {
 
     @Test
     public void containerIsRecreatedIfMemoryChanged() {
-        NodeSpec.Builder specBuilder = nodeBuilder
-                .state(NodeState.active)
+        NodeSpec.Builder specBuilder = nodeBuilder(NodeState.active)
                 .wantedDockerImage(dockerImage).currentDockerImage(dockerImage)
                 .wantedVespaVersion(vespaVersion).currentVespaVersion(vespaVersion)
                 .wantedRestartGeneration(2).currentRestartGeneration(1);
@@ -304,8 +285,7 @@ public class NodeAgentImplTest {
 
     @Test
     public void noRestartIfOrchestratorSuspendFails() {
-        final NodeSpec node = nodeBuilder
-                .state(NodeState.active)
+        final NodeSpec node = nodeBuilder(NodeState.active)
                 .wantedDockerImage(dockerImage).currentDockerImage(dockerImage)
                 .wantedVespaVersion(vespaVersion).currentVespaVersion(vespaVersion)
                 .wantedRestartGeneration(2).currentRestartGeneration(1)
@@ -332,11 +312,9 @@ public class NodeAgentImplTest {
     @Test
     public void recreatesContainerIfRebootWanted() {
         final long wantedRebootGeneration = 2;
-        final NodeSpec node = nodeBuilder
-                .state(NodeState.active)
+        final NodeSpec node = nodeBuilder(NodeState.active)
                 .wantedDockerImage(dockerImage).currentDockerImage(dockerImage)
                 .wantedVespaVersion(vespaVersion).currentVespaVersion(vespaVersion)
-                .wantedRestartGeneration(1).currentRestartGeneration(1)
                 .wantedRebootGeneration(wantedRebootGeneration).currentRebootGeneration(1)
                 .build();
 
@@ -373,8 +351,7 @@ public class NodeAgentImplTest {
 
     @Test
     public void failedNodeRunningContainerShouldStillBeRunning() {
-        final NodeSpec node = nodeBuilder
-                .state(NodeState.failed)
+        final NodeSpec node = nodeBuilder(NodeState.failed)
                 .wantedDockerImage(dockerImage).currentDockerImage(dockerImage)
                 .wantedVespaVersion(vespaVersion).currentVespaVersion(vespaVersion)
                 .build();
@@ -393,9 +370,7 @@ public class NodeAgentImplTest {
 
     @Test
     public void readyNodeLeadsToNoAction() {
-        final NodeSpec node = nodeBuilder
-                .state(NodeState.ready)
-                .build();
+        final NodeSpec node = nodeBuilder(NodeState.ready).build();
 
         NodeAgentContext context = createContext(node);
         NodeAgentImpl nodeAgent = makeNodeAgent(null,false);
@@ -417,8 +392,7 @@ public class NodeAgentImplTest {
 
     @Test
     public void inactiveNodeRunningContainerShouldStillBeRunning() {
-        final NodeSpec node = nodeBuilder
-                .state(NodeState.inactive)
+        final NodeSpec node = nodeBuilder(NodeState.inactive)
                 .wantedDockerImage(dockerImage).currentDockerImage(dockerImage)
                 .wantedVespaVersion(vespaVersion).currentVespaVersion(vespaVersion)
                 .build();
@@ -439,8 +413,7 @@ public class NodeAgentImplTest {
 
     @Test
     public void reservedNodeDoesNotUpdateNodeRepoWithVersion() {
-        final NodeSpec node = nodeBuilder
-                .state(NodeState.reserved)
+        final NodeSpec node = nodeBuilder(NodeState.reserved)
                 .wantedDockerImage(dockerImage)
                 .wantedVespaVersion(vespaVersion)
                 .build();
@@ -456,14 +429,12 @@ public class NodeAgentImplTest {
     }
 
     private void nodeRunningContainerIsTakenDownAndCleanedAndRecycled(NodeState nodeState, Optional<Long> wantedRestartGeneration) {
-        wantedRestartGeneration.ifPresent(restartGeneration -> nodeBuilder
+        NodeSpec.Builder builder = nodeBuilder(nodeState)
+                .wantedDockerImage(dockerImage).currentDockerImage(dockerImage);
+        wantedRestartGeneration.ifPresent(restartGeneration -> builder
                 .wantedRestartGeneration(restartGeneration).currentRestartGeneration(restartGeneration));
 
-        final NodeSpec node = nodeBuilder
-                .state(nodeState)
-                .wantedDockerImage(dockerImage).currentDockerImage(dockerImage)
-                .build();
-
+        NodeSpec node = builder.build();
         NodeAgentContext context = createContext(node);
         NodeAgentImpl nodeAgent = makeNodeAgent(dockerImage, true);
 
@@ -501,9 +472,8 @@ public class NodeAgentImplTest {
 
     @Test
     public void provisionedNodeIsMarkedAsDirty() {
-        final NodeSpec node = nodeBuilder
+        final NodeSpec node = nodeBuilder(NodeState.provisioned)
                 .wantedDockerImage(dockerImage)
-                .state(NodeState.provisioned)
                 .build();
 
         NodeAgentContext context = createContext(node);
@@ -516,8 +486,7 @@ public class NodeAgentImplTest {
 
     @Test
     public void testRestartDeadContainerAfterNodeAdminRestart() {
-        final NodeSpec node = nodeBuilder
-                .state(NodeState.active)
+        final NodeSpec node = nodeBuilder(NodeState.active)
                 .currentDockerImage(dockerImage).wantedDockerImage(dockerImage)
                 .currentVespaVersion(vespaVersion)
                 .build();
@@ -537,8 +506,7 @@ public class NodeAgentImplTest {
 
     @Test
     public void resumeProgramRunsUntilSuccess() {
-        final NodeSpec node = nodeBuilder
-                .state(NodeState.active)
+        final NodeSpec node = nodeBuilder(NodeState.active)
                 .wantedDockerImage(dockerImage).currentDockerImage(dockerImage)
                 .currentVespaVersion(vespaVersion)
                 .wantedRestartGeneration(1).currentRestartGeneration(1)
@@ -575,8 +543,7 @@ public class NodeAgentImplTest {
 
     @Test
     public void start_container_subtask_failure_leads_to_container_restart() {
-        final NodeSpec node = nodeBuilder
-                .state(NodeState.active)
+        final NodeSpec node = nodeBuilder(NodeState.active)
                 .wantedDockerImage(dockerImage)
                 .wantedVespaVersion(vespaVersion)
                 .wantedRestartGeneration(1).currentRestartGeneration(1)
@@ -612,8 +579,7 @@ public class NodeAgentImplTest {
 
     @Test
     public void testRunningConfigServer() {
-        final NodeSpec node = nodeBuilder
-                .state(NodeState.active)
+        final NodeSpec node = nodeBuilder(NodeState.active)
                 .type(NodeType.config)
                 .wantedDockerImage(dockerImage)
                 .wantedVespaVersion(vespaVersion)
@@ -656,11 +622,9 @@ public class NodeAgentImplTest {
 
     @Test
     public void initial_cpu_cap_test() {
-        NodeSpec.Builder specBuilder = nodeBuilder
-                .state(NodeState.active)
+        NodeSpec.Builder specBuilder = nodeBuilder(NodeState.active)
                 .wantedDockerImage(dockerImage).currentDockerImage(dockerImage)
-                .wantedVespaVersion(vespaVersion).currentVespaVersion(vespaVersion)
-                .wantedRestartGeneration(1).currentRestartGeneration(1);
+                .wantedVespaVersion(vespaVersion).currentVespaVersion(vespaVersion);
 
         NodeAgentContext context = createContext(specBuilder.build());
         NodeAgentImpl nodeAgent = makeNodeAgent(null, false, Duration.ofSeconds(30));
@@ -712,8 +676,7 @@ public class NodeAgentImplTest {
 
     @Test
     public void resumes_normally_if_container_is_already_capped_on_start() {
-        NodeSpec.Builder specBuilder = nodeBuilder
-                .state(NodeState.active)
+        NodeSpec.Builder specBuilder = nodeBuilder(NodeState.active)
                 .wantedDockerImage(dockerImage).currentDockerImage(dockerImage)
                 .wantedVespaVersion(vespaVersion).currentVespaVersion(vespaVersion)
                 .wantedRestartGeneration(1).currentRestartGeneration(1);
@@ -814,4 +777,8 @@ public class NodeAgentImplTest {
     private NodeAgentContext createContext(NodeSpec nodeSpec) {
         return new NodeAgentContextImpl.Builder(nodeSpec).build();
     }
+
+    private NodeSpec.Builder nodeBuilder(NodeState state) {
+        return NodeSpec.Builder.testSpec(hostName, state).resources(resources);
+    }
 }