Remove Mockito usage from MetricsReporterTest

1 files changed, 19 insertions, 9 deletions

diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/chef/ChefMock.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/chef/ChefMock.java
index 1b2dad34b8d..bd19cfe6ce1 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/chef/ChefMock.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/chef/ChefMock.java
@@ -16,7 +16,6 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.stream.Collectors;
 
 /**
  * @author mpolden
@@ -24,11 +23,14 @@ import java.util.stream.Collectors;
 public class ChefMock implements Chef {
 
     private final NodeResult result;
+    private final PartialNodeResult partialResult;
     private final List<String> chefEnvironments;
 
     public ChefMock() {
         result = new NodeResult();
         result.rows = new ArrayList<>();
+        partialResult = new PartialNodeResult();
+        partialResult.rows = new ArrayList<>();
         chefEnvironments = new ArrayList<>();
         chefEnvironments.add("hosted-verified-prod");
         chefEnvironments.add("hosted-infra-cd");
@@ -59,8 +61,14 @@ public class ChefMock implements Chef {
         return null;
     }
 
-    public void addSearchResult(ChefNode node) {
+    public ChefMock addSearchResult(ChefNode node) {
         result.rows.add(node);
+        return this;
+    }
+
+    public ChefMock addPartialResult(List<PartialNode> partialNodes) {
+        partialResult.rows.addAll(partialNodes);
+        return this;
     }
 
     @Override
@@ -76,13 +84,15 @@ public class ChefMock implements Chef {
     @Override
     public PartialNodeResult partialSearchNodes(String query, List<AttributeMapping> returnAttributes) {
         PartialNodeResult partialNodeResult = new PartialNodeResult();
-        partialNodeResult.rows = result.rows.stream()
-                .map(chefNode -> {
-                    Map<String, String> data = new HashMap<>();
-                    data.put("fqdn", chefNode.name);
-                    return new PartialNode(data);
-                })
-                .collect(Collectors.toList());
+        partialNodeResult.rows = new ArrayList<>();
+        partialNodeResult.rows.addAll(partialResult.rows);
+        result.rows.stream()
+                   .map(chefNode -> {
+                       Map<String, String> data = new HashMap<>();
+                       data.put("fqdn", chefNode.name);
+                       return new PartialNode(data);
+                   })
+                   .forEach(node -> partialNodeResult.rows.add(node));
         return partialNodeResult;
     }