1 files changed, 31 insertions, 7 deletions

diff --git a/catalog-be/src/test/java/org/openecomp/sdc/be/components/utils/ComponentInstanceBuilder.java b/catalog-be/src/test/java/org/openecomp/sdc/be/components/utils/ComponentInstanceBuilder.java
index a99f08c20d..f461806010 100644
--- a/catalog-be/src/test/java/org/openecomp/sdc/be/components/utils/ComponentInstanceBuilder.java
+++ b/catalog-be/src/test/java/org/openecomp/sdc/be/components/utils/ComponentInstanceBuilder.java
@@ -1,17 +1,25 @@
 package org.openecomp.sdc.be.components.utils;
 
-import java.util.HashMap;
-import java.util.Map;
-
 import org.openecomp.sdc.be.model.ArtifactDefinition;
+import org.openecomp.sdc.be.model.CapabilityDefinition;
 import org.openecomp.sdc.be.model.ComponentInstance;
 
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.stream.Stream;
+
 public class ComponentInstanceBuilder {
 
     private ComponentInstance componentInstance;
 
     public ComponentInstanceBuilder() {
         componentInstance = new ComponentInstance();
+        componentInstance.setCapabilities(new HashMap<>());
+        componentInstance.setDeploymentArtifacts(new HashMap<>());
+    }
+
+    public ComponentInstanceBuilder(ComponentInstance componentInstance) {
+        this.componentInstance = componentInstance;
     }
 
     public ComponentInstanceBuilder setName(String name) {
@@ -19,6 +27,16 @@ public class ComponentInstanceBuilder {
         return this;
     }
 
+    public ComponentInstanceBuilder setNormalizedName(String name) {
+        componentInstance.setNormalizedName(name);
+        return this;
+    }
+
+    public ComponentInstanceBuilder setUniqueId(String uniqueId) {
+        componentInstance.setUniqueId(uniqueId);
+        return this;
+    }
+
     public ComponentInstanceBuilder setComponentUid(String componentUid) {
         componentInstance.setComponentUid(componentUid);
         return this;
@@ -35,14 +53,20 @@ public class ComponentInstanceBuilder {
     }
 
     public ComponentInstanceBuilder addDeploymentArtifact(ArtifactDefinition artifactDefinition) {
-        Map<String, ArtifactDefinition> deploymentArtifacts = componentInstance.getDeploymentArtifacts();
-        if (deploymentArtifacts == null) {
-            componentInstance.setDeploymentArtifacts(new HashMap<>());
-        }
         componentInstance.getDeploymentArtifacts().put(artifactDefinition.getArtifactName(), artifactDefinition);
         return this;
     }
 
+    public ComponentInstanceBuilder addCapability(CapabilityDefinition capabilityDefinition) {
+        componentInstance.getCapabilities().computeIfAbsent(capabilityDefinition.getType(), key -> new ArrayList<>()).add(capabilityDefinition);
+        return this;
+    }
+
+    public ComponentInstanceBuilder addCapabilities(CapabilityDefinition ... capabilities) {
+        Stream.of(capabilities).forEach(this::addCapability);
+        return this;
+    }
+
     public ComponentInstance build() {
         return componentInstance;
     }