Fix error adding capability to VFC

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

1 files changed, 10 insertions, 6 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() {