From c3e180db7928d9eb78b565e16749542a6a663d3d Mon Sep 17 00:00:00 2001
From: Miguel Caballer <micafer1@upv.es>
Date: Tue, 2 May 2017 13:34:00 +0200
Subject: [PATCH] Fix error getting relationshps in case of custom_def
 capability

Change-Id: I5bbc21113451d17e887ee79426e4205d40f57f6a
Related-Bug: #1687598
---
 toscaparser/elements/nodetype.py              | 13 +++++++----
 .../tests/data/custom_types/custom_cap.yaml   | 22 ++++++++++++++++++
 .../tests/data/test_custom_capabilty.yaml     | 23 +++++++++++++++++++
 toscaparser/tests/test_toscatpl.py            |  6 +++++
 4 files changed, 60 insertions(+), 4 deletions(-)
 create mode 100644 toscaparser/tests/data/custom_types/custom_cap.yaml
 create mode 100644 toscaparser/tests/data/test_custom_capabilty.yaml

diff --git a/toscaparser/elements/nodetype.py b/toscaparser/elements/nodetype.py
index 57f40e1..4889ed8 100644
--- a/toscaparser/elements/nodetype.py
+++ b/toscaparser/elements/nodetype.py
@@ -83,7 +83,6 @@ class NodeType(StatefulEntityType):
                                 captype = value['capability']
                                 value = (self.
                                          _get_node_type_by_cap(captype))
-                            relation = self._get_relation(key, value)
                             keyword = key
                             node_type = value
                 rtype = RelationshipType(relation, keyword, self.custom_def)
@@ -102,9 +101,15 @@ class NodeType(StatefulEntityType):
         node_types = [node_type for node_type in self.TOSCA_DEF.keys()
                       if node_type.startswith(self.NODE_PREFIX) and
                       node_type != 'tosca.nodes.Root']
+        custom_node_types = [node_type for node_type in self.custom_def.keys()
+                             if node_type.startswith(self.NODE_PREFIX) and
+                             node_type != 'tosca.nodes.Root']
 
-        for node_type in node_types:
-            node_def = self.TOSCA_DEF[node_type]
+        for node_type in node_types + custom_node_types:
+            if node_type in self.TOSCA_DEF:
+                node_def = self.TOSCA_DEF[node_type]
+            else:
+                node_def = self.custom_def[node_type]
             if isinstance(node_def, dict) and 'capabilities' in node_def:
                 node_caps = node_def['capabilities']
                 for value in node_caps.values():
@@ -114,7 +119,7 @@ class NodeType(StatefulEntityType):
 
     def _get_relation(self, key, ndtype):
         relation = None
-        ntype = NodeType(ndtype)
+        ntype = NodeType(ndtype, self.custom_def)
         caps = ntype.get_capabilities()
         if caps and key in caps.keys():
             c = caps[key]
diff --git a/toscaparser/tests/data/custom_types/custom_cap.yaml b/toscaparser/tests/data/custom_types/custom_cap.yaml
new file mode 100644
index 0000000..018bcf6
--- /dev/null
+++ b/toscaparser/tests/data/custom_types/custom_cap.yaml
@@ -0,0 +1,22 @@
+tosca_definitions_version: tosca_simple_yaml_1_0
+
+capability_types:
+
+  tosca.capabilities.SomeCap:
+    derived_from: tosca.capabilities.Container
+
+node_types:
+
+  tosca.nodes.NodeWithReq:
+    derived_from: tosca.nodes.SoftwareComponent
+    requirements:
+      - host:
+          capability: tosca.capabilities.SomeCap
+          relationship: tosca.relationships.HostedOn
+          occurrences: [1, 1]
+
+  tosca.nodes.NodeWithCap:
+    derived_from: tosca.nodes.SoftwareComponent
+    capabilities:
+      host:
+        type: tosca.capabilities.SomeCap
diff --git a/toscaparser/tests/data/test_custom_capabilty.yaml b/toscaparser/tests/data/test_custom_capabilty.yaml
new file mode 100644
index 0000000..03a8a07
--- /dev/null
+++ b/toscaparser/tests/data/test_custom_capabilty.yaml
@@ -0,0 +1,23 @@
+tosca_definitions_version: tosca_simple_yaml_1_0
+
+description: TOSCA simple profile to test a custom defined capability
+
+imports:
+  - custom_types/custom_cap.yaml
+
+topology_template:
+
+  node_templates:
+
+    node_req:
+      type: tosca.nodes.NodeWithReq
+      requirements:
+        - host: node_cap
+
+    node_cap:
+      type: tosca.nodes.NodeWithCap
+      requirements:
+        - host: server
+
+    server:
+      type: tosca.nodes.Compute
diff --git a/toscaparser/tests/test_toscatpl.py b/toscaparser/tests/test_toscatpl.py
index 77232df..88d8d6e 100644
--- a/toscaparser/tests/test_toscatpl.py
+++ b/toscaparser/tests/test_toscatpl.py
@@ -851,3 +851,9 @@ class ToscaTemplateTest(TestCase):
         self.assertEqual(
             ['ALRM1', 'SP1', 'SP2'],
             sorted([policy.name for policy in tosca.policies]))
+
+    def test_custom_capability(self):
+        tosca_tpl = os.path.join(
+            os.path.dirname(os.path.abspath(__file__)),
+            "data/test_custom_capabilty.yaml")
+        ToscaTemplate(tosca_tpl)