Make it easier to modify NodeSpec in NodeAgentContext.Builder

3 files changed, 5 insertions, 14 deletions

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 ef8ea60bee3..e20480e14ef 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
@@ -18,6 +18,7 @@ import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.function.Function;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -183,8 +184,8 @@ public class NodeAgentContextImpl implements NodeAgentContext {
                     .flavor("d-2-8-50");
         }
 
-        public Builder nodeType(NodeType nodeType) {
-            this.nodeSpecBuilder.type(nodeType);
+        public Builder nodeSpecBuilder(Function<NodeSpec.Builder, NodeSpec.Builder> nodeSpecBuilderModifier) {
+            this.nodeSpecBuilder = nodeSpecBuilderModifier.apply(nodeSpecBuilder);
             return this;
         }
 
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImpl.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImpl.java
index 98583262136..b7e7b97cdd8 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImpl.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImpl.java
@@ -355,7 +355,7 @@ public class NodeAgentImpl implements NodeAgent {
     }
 
     private boolean noCpuCap(ZoneApi zone) {
-        return zone.getEnvironment() == Environment.dev || (zone.getSystemName().isCd());
+        return zone.getEnvironment() == Environment.dev || zone.getSystemName().isCd();
     }
 
     private boolean downloadImageIfNeeded(NodeSpec node, Optional<Container> container) {
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 ca9b05a3ff6..b33f52ff629 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,11 +1,8 @@
 // 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.test.ManualClock;
 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.nodeagent.NodeAgentContext;
 import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContextImpl;
 import org.junit.Test;
@@ -157,14 +154,7 @@ public class NodeAdminImplTest {
     }
 
     private NodeAgentContext createNodeAgentContext(String hostname) {
-        NodeSpec nodeSpec = new NodeSpec.Builder()
-                .hostname(hostname)
-                .state(NodeState.active)
-                .type(NodeType.tenant)
-                .flavor("default")
-                .build();
-
-        return new NodeAgentContextImpl.Builder(nodeSpec).build();
+        return new NodeAgentContextImpl.Builder(hostname).build();
     }
 
     private NodeAgentWithScheduler mockNodeAgentWithSchedulerFactory(NodeAgentContext context) {