Revert "Revert "Move node operations into a separate class""

4 files changed, 21 insertions, 7 deletions

diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/identitydocument/IdentityDocumentGenerator.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/identitydocument/IdentityDocumentGenerator.java
index 123a9bc7449..cf4ca85ecfd 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/identitydocument/IdentityDocumentGenerator.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/identitydocument/IdentityDocumentGenerator.java
@@ -47,7 +47,7 @@ public class IdentityDocumentGenerator {
 
     public SignedIdentityDocument generateSignedIdentityDocument(String hostname, IdentityType identityType) {
         try {
-            Node node = nodeRepository.getNode(hostname).orElseThrow(() -> new RuntimeException("Unable to find node " + hostname));
+            Node node = nodeRepository.nodes().getNode(hostname).orElseThrow(() -> new RuntimeException("Unable to find node " + hostname));
             Allocation allocation = node.allocation().orElseThrow(() -> new RuntimeException("No allocation for node " + node.hostname()));
             VespaUniqueInstanceId providerUniqueId = new VespaUniqueInstanceId(
                     allocation.membership().index(),
diff --git a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/instanceconfirmation/InstanceValidator.java b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/instanceconfirmation/InstanceValidator.java
index c258e1be466..40782455025 100644
--- a/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/instanceconfirmation/InstanceValidator.java
+++ b/athenz-identity-provider-service/src/main/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/instanceconfirmation/InstanceValidator.java
@@ -137,7 +137,7 @@ public class InstanceValidator {
         }
 
         // Find node matching vespa unique id
-        Node node = nodeRepository.getNodes().stream()
+        Node node = nodeRepository.nodes().getNodes().stream()
                 .filter(n -> n.allocation().isPresent())
                 .filter(n -> nodeMatchesVespaUniqueId(n, vespaUniqueInstanceId))
                 .findFirst() // Should be only one
diff --git a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/identitydocument/IdentityDocumentGeneratorTest.java b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/identitydocument/IdentityDocumentGeneratorTest.java
index 0a757b5b17b..a6dfc6e9b9e 100644
--- a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/identitydocument/IdentityDocumentGeneratorTest.java
+++ b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/identitydocument/IdentityDocumentGeneratorTest.java
@@ -24,6 +24,7 @@ import com.yahoo.vespa.hosted.provision.NodeRepository;
 import com.yahoo.vespa.hosted.provision.node.Allocation;
 import com.yahoo.vespa.hosted.provision.node.Generation;
 import com.yahoo.vespa.hosted.provision.node.IP;
+import com.yahoo.vespa.hosted.provision.node.Nodes;
 import com.yahoo.vespa.hosted.provision.testutils.MockNodeFlavors;
 import org.junit.Test;
 
@@ -68,8 +69,11 @@ public class IdentityDocumentGeneratorTest {
                                                    NodeType.tenant)
                 .allocation(allocation).build();
         NodeRepository nodeRepository = mock(NodeRepository.class);
-        when(nodeRepository.getNode(eq(parentHostname))).thenReturn(Optional.of(parentNode));
-        when(nodeRepository.getNode(eq(containerHostname))).thenReturn(Optional.of(containerNode));
+        Nodes nodes = mock(Nodes.class);
+        when(nodeRepository.nodes()).thenReturn(nodes);
+
+        when(nodes.getNode(eq(parentHostname))).thenReturn(Optional.of(parentNode));
+        when(nodes.getNode(eq(containerHostname))).thenReturn(Optional.of(containerNode));
         AutoGeneratedKeyProvider keyProvider = new AutoGeneratedKeyProvider();
 
         String dnsSuffix = "vespa.dns.suffix";
diff --git a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/instanceconfirmation/InstanceValidatorTest.java b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/instanceconfirmation/InstanceValidatorTest.java
index c5dbd9dd7b2..80bfda811d0 100644
--- a/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/instanceconfirmation/InstanceValidatorTest.java
+++ b/athenz-identity-provider-service/src/test/java/com/yahoo/vespa/hosted/athenz/instanceproviderservice/instanceconfirmation/InstanceValidatorTest.java
@@ -23,6 +23,7 @@ import com.yahoo.vespa.hosted.athenz.instanceproviderservice.KeyProvider;
 import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
 import com.yahoo.vespa.hosted.provision.node.IP;
+import com.yahoo.vespa.hosted.provision.node.Nodes;
 import com.yahoo.vespa.hosted.provision.testutils.MockNodeFlavors;
 import org.junit.Test;
 
@@ -119,12 +120,14 @@ public class InstanceValidatorTest {
     @Test
     public void accepts_valid_refresh_requests() {
         NodeRepository nodeRepository = mock(NodeRepository.class);
+        Nodes nodes = mock(Nodes.class);
+        when(nodeRepository.nodes()).thenReturn(nodes);
         InstanceValidator instanceValidator = new InstanceValidator(null, null, nodeRepository, new IdentityDocumentSigner(), vespaTenantDomain);
 
         List<Node> nodeList = createNodes(10);
         Node node = nodeList.get(0);
         nodeList = allocateNode(nodeList, node, applicationId);
-        when(nodeRepository.getNodes()).thenReturn(nodeList);
+        when(nodes.getNodes()).thenReturn(nodeList);
         String nodeIp = node.ipConfig().primary().stream().findAny().orElseThrow(() -> new RuntimeException("No ipaddress for mocked node"));
         InstanceConfirmation instanceConfirmation = createRefreshInstanceConfirmation(applicationId, domain, service, ImmutableList.of(nodeIp));
 
@@ -134,12 +137,15 @@ public class InstanceValidatorTest {
     @Test
     public void rejects_refresh_on_ip_mismatch() {
         NodeRepository nodeRepository = mock(NodeRepository.class);
+        Nodes nodes = mock(Nodes.class);
+        when(nodeRepository.nodes()).thenReturn(nodes);
+
         InstanceValidator instanceValidator = new InstanceValidator(null, null, nodeRepository, new IdentityDocumentSigner(), vespaTenantDomain);
 
         List<Node> nodeList = createNodes(10);
         Node node = nodeList.get(0);
         nodeList = allocateNode(nodeList, node, applicationId);
-        when(nodeRepository.getNodes()).thenReturn(nodeList);
+        when(nodes.getNodes()).thenReturn(nodeList);
         String nodeIp = node.ipConfig().primary().stream().findAny().orElseThrow(() -> new RuntimeException("No ipaddress for mocked node"));
 
         // Add invalid ip to list of ip addresses
@@ -151,10 +157,14 @@ public class InstanceValidatorTest {
     @Test
     public void rejects_refresh_when_node_is_not_allocated() {
         NodeRepository nodeRepository = mock(NodeRepository.class);
+        Nodes nodes = mock(Nodes.class);
+        when(nodeRepository.nodes()).thenReturn(nodes);
+
         InstanceValidator instanceValidator = new InstanceValidator(null, null, nodeRepository, new IdentityDocumentSigner(), vespaTenantDomain);
 
         List<Node> nodeList = createNodes(10);
-        when(nodeRepository.getNodes()).thenReturn(nodeList);
+
+        when(nodes.getNodes()).thenReturn(nodeList);
         InstanceConfirmation instanceConfirmation = createRefreshInstanceConfirmation(applicationId, domain, service, ImmutableList.of("::11"));
 
         assertFalse(instanceValidator.isValidRefresh(instanceConfirmation));