Fix error adding capability to VFC

Signed-off-by: MichaelMorris <michael.morris@est.tech>
Issue-ID: SDC-4098
Change-Id: Id3c71787b9c11bbf4316a9a23f01199a009f2b49

3 files changed, 45 insertions, 11 deletions

diff --git a/catalog-model/src/main/java/org/openecomp/sdc/be/model/jsonjanusgraph/operations/CapabilitiesOperation.java b/catalog-model/src/main/java/org/openecomp/sdc/be/model/jsonjanusgraph/operations/CapabilitiesOperation.java
index 65673345e7..5977acae1e 100644
--- a/catalog-model/src/main/java/org/openecomp/sdc/be/model/jsonjanusgraph/operations/CapabilitiesOperation.java
+++ b/catalog-model/src/main/java/org/openecomp/sdc/be/model/jsonjanusgraph/operations/CapabilitiesOperation.java
@@ -92,7 +92,7 @@ public class CapabilitiesOperation extends BaseOperation {
         }
     }
 
-    private StorageOperationStatus createOrUpdateCapabilityProperties(String componentId, TopologyTemplate toscaElement,
+    private StorageOperationStatus createOrUpdateCapabilityProperties(String componentId, ToscaElement toscaElement,
                                                                       Map<String, MapPropertiesDataDefinition> propertiesMap) {
         GraphVertex toscaElementV = janusGraphDao.getVertexById(componentId, JsonParseFlagEnum.NoParse).left().on(this::throwStorageException);
         Map<String, MapPropertiesDataDefinition> capabilitiesProperties = toscaElement.getCapabilitiesProperties();
@@ -116,17 +116,21 @@ public class CapabilitiesOperation extends BaseOperation {
             .updateFullToscaData(toscaElementV, EdgeLabelEnum.CAPABILITIES_PROPERTIES, VertexTypeEnum.CAPABILITIES_PROPERTIES, propertiesMap);
     }
 
-    public StorageOperationStatus createOrUpdateCapabilityProperties(String componentId, Map<String, MapPropertiesDataDefinition> propertiesMap) {
+    public StorageOperationStatus createOrUpdateCapabilityProperties(String componentId, boolean isTopologyTemplate, Map<String, MapPropertiesDataDefinition> propertiesMap) {
         StorageOperationStatus propertiesStatusRes = null;
         if (MapUtils.isNotEmpty(propertiesMap)) {
-            propertiesStatusRes = createOrUpdateCapabilityProperties(componentId, getTopologyTemplate(componentId), propertiesMap);
+            propertiesStatusRes = createOrUpdateCapabilityProperties(componentId, getToscaElement(componentId, isTopologyTemplate), propertiesMap);
         }
         return propertiesStatusRes;
     }
 
-    private TopologyTemplate getTopologyTemplate(String componentId) {
-        return (TopologyTemplate) topologyTemplateOperation.getToscaElement(componentId, getFilterComponentWithCapProperties()).left()
-            .on(this::throwStorageException);
+    private ToscaElement getToscaElement(String componentId, boolean isTopologyTemplate) {
+        if (isTopologyTemplate){
+            return topologyTemplateOperation.getToscaElement(componentId, getFilterComponentWithCapProperties()).left()
+                .on(this::throwStorageException);
+        }
+        return nodeTypeOperation.getToscaElement(componentId, getFilterComponentWithCapProperties()).left()
+                .on(this::throwStorageException);
     }
 
     private ComponentParametersView getFilterComponentWithCapProperties() {
diff --git a/catalog-model/src/test/java/org/openecomp/sdc/be/model/jsonjanusgraph/operations/CapabilitiesOperationTest.java b/catalog-model/src/test/java/org/openecomp/sdc/be/model/jsonjanusgraph/operations/CapabilitiesOperationTest.java
index 6b5a210d7a..0edde3a0f6 100644
--- a/catalog-model/src/test/java/org/openecomp/sdc/be/model/jsonjanusgraph/operations/CapabilitiesOperationTest.java
+++ b/catalog-model/src/test/java/org/openecomp/sdc/be/model/jsonjanusgraph/operations/CapabilitiesOperationTest.java
@@ -28,6 +28,7 @@ import org.openecomp.sdc.be.dao.janusgraph.JanusGraphDao;
 import org.openecomp.sdc.be.dao.janusgraph.JanusGraphOperationStatus;
 import org.openecomp.sdc.be.datatypes.elements.MapPropertiesDataDefinition;
 import org.openecomp.sdc.be.datatypes.elements.PropertyDataDefinition;
+import org.openecomp.sdc.be.model.jsonjanusgraph.datamodel.NodeType;
 import org.openecomp.sdc.be.model.jsonjanusgraph.datamodel.TopologyTemplate;
 import org.openecomp.sdc.be.model.jsonjanusgraph.utils.CapabilityTestUtils;
 import org.openecomp.sdc.be.model.operations.api.StorageOperationStatus;
@@ -48,6 +49,9 @@ public class CapabilitiesOperationTest {
     private  JanusGraphDao mockJanusGraphDao;
     @Mock
     private TopologyTemplateOperation topologyTemplateOperation;
+    @Mock
+    private NodeTypeOperation nodeTypeOperation;
+
 
     @Before
     public void setUp() {
@@ -64,10 +68,36 @@ public class CapabilitiesOperationTest {
         topologyTemplate.setCapabilitiesProperties(capPropsForTopologyTemplate);
 
         when(topologyTemplateOperation.getToscaElement(anyString(), any())).thenReturn(Either.left(topologyTemplate));
+
+        NodeType nodeType = new NodeType();
+        Map<String, MapPropertiesDataDefinition> capPropsForNodeType = CapabilityTestUtils
+                .createCapPropsForTopologyTemplate(nodeType);
+        nodeType.setCapabilitiesProperties(capPropsForNodeType);
+
+        when(nodeTypeOperation.getToscaElement(anyString(), any())).thenReturn(Either.left(nodeType));
     }
 
     @Test
-    public void testCreateOrUpdateCapabilitiesProperties() {
+    public void testCreateOrUpdateCapabilitiesPropertiesTopologyTemplate() {
+
+        Map<String, PropertyDataDefinition> mapToscaDataDefinition = new HashMap<>();
+        PropertyDataDefinition propertyDataDefinition = new PropertyDataDefinition();
+        propertyDataDefinition.setUniqueId("ComponentInput1_uniqueId");
+        propertyDataDefinition.setName("propName");
+        mapToscaDataDefinition.put(propertyDataDefinition.getUniqueId(), propertyDataDefinition);
+        MapPropertiesDataDefinition  mapPropertiesDataDefinition = new MapPropertiesDataDefinition(mapToscaDataDefinition);
+
+        Map<String, MapPropertiesDataDefinition> propertiesMap = new HashMap<>();
+        propertiesMap.put(propertyDataDefinition.getUniqueId(), mapPropertiesDataDefinition);
+
+        StorageOperationStatus operationStatus = operation.createOrUpdateCapabilityProperties("componentId", true,
+                propertiesMap);
+
+        Assert.assertEquals(StorageOperationStatus.OK, operationStatus);
+    }
+    
+    @Test
+    public void testCreateOrUpdateCapabilitiesPropertiesToscaTemplate() {
 
         Map<String, PropertyDataDefinition> mapToscaDataDefinition = new HashMap<>();
         PropertyDataDefinition propertyDataDefinition = new PropertyDataDefinition();
@@ -79,9 +109,9 @@ public class CapabilitiesOperationTest {
         Map<String, MapPropertiesDataDefinition> propertiesMap = new HashMap<>();
         propertiesMap.put(propertyDataDefinition.getUniqueId(), mapPropertiesDataDefinition);
 
-        StorageOperationStatus operationStatus = operation.createOrUpdateCapabilityProperties("componentId",
+        StorageOperationStatus operationStatus = operation.createOrUpdateCapabilityProperties("componentId", false,
                 propertiesMap);
 
         Assert.assertEquals(StorageOperationStatus.OK, operationStatus);
     }
-}
\ No newline at end of file
+}
diff --git a/catalog-model/src/test/java/org/openecomp/sdc/be/model/jsonjanusgraph/utils/CapabilityTestUtils.java b/catalog-model/src/test/java/org/openecomp/sdc/be/model/jsonjanusgraph/utils/CapabilityTestUtils.java
index 8a17de4054..ed2d43a8ed 100644
--- a/catalog-model/src/test/java/org/openecomp/sdc/be/model/jsonjanusgraph/utils/CapabilityTestUtils.java
+++ b/catalog-model/src/test/java/org/openecomp/sdc/be/model/jsonjanusgraph/utils/CapabilityTestUtils.java
@@ -24,7 +24,7 @@ import org.openecomp.sdc.be.datatypes.elements.SchemaDefinition;
 import org.openecomp.sdc.be.model.CapabilityDefinition;
 import org.openecomp.sdc.be.model.ComponentInstanceProperty;
 import org.openecomp.sdc.be.model.jsonjanusgraph.datamodel.TopologyTemplate;
-
+import org.openecomp.sdc.be.model.jsonjanusgraph.datamodel.ToscaElement;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -32,7 +32,7 @@ import java.util.Map;
 
 public class CapabilityTestUtils {
 
-    public static Map<String, MapPropertiesDataDefinition> createCapPropsForTopologyTemplate(TopologyTemplate topologyTemplate) {
+    public static Map<String, MapPropertiesDataDefinition> createCapPropsForTopologyTemplate(ToscaElement topologyTemplate) {
         Map<String, ListCapabilityDataDefinition> capabilitiesMap = new HashMap<>();
 
         List<CapabilityDataDefinition> capabilityDefinitions = new ArrayList<>();