Follow API change

2 files changed, 19 insertions, 5 deletions

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));