diff --git a/netdiff/parsers/netjson.py b/netdiff/parsers/netjson.py index bc09bf4..1544020 100644 --- a/netdiff/parsers/netjson.py +++ b/netdiff/parsers/netjson.py @@ -32,6 +32,7 @@ def parse(self, data): # create graph for node in data['nodes']: graph.add_node(node['id'], + label=node['label'] if 'label' in node else None, local_addresses=node.get('local_addresses', []), **node.get('properties', {})) for link in data['links']: diff --git a/netdiff/utils.py b/netdiff/utils.py index 5ffc69f..8b71b24 100644 --- a/netdiff/utils.py +++ b/netdiff/utils.py @@ -136,12 +136,16 @@ def _netjson_networkgraph(protocol, version, revision, metric, # must copy properties dict to avoid modifying data properties = node[1].copy() local_addresses = properties.pop('local_addresses', None) + label = properties.pop('label', None) # append local_addresses only if not empty if local_addresses: netjson_node['local_addresses'] = local_addresses # append properties only if not empty if properties: netjson_node['properties'] = properties + # append label only if not empty + if label: + netjson_node['label'] = label node_list.append(netjson_node) # prepare links link_list = [] diff --git a/tests/static/netjson-2-links.json b/tests/static/netjson-2-links.json index e55df7b..d2b3fbb 100644 --- a/tests/static/netjson-2-links.json +++ b/tests/static/netjson-2-links.json @@ -10,6 +10,7 @@ "local_addresses": [ "192.168.1.3" ], + "label": "nodeA", "properties": { "hostname": "router.3nnx" } @@ -19,6 +20,7 @@ "local_addresses": [ "192.168.1.2" ], + "label": "nodeB", "properties": { "hostname": "router2.nnx" } diff --git a/tests/test_netjson.py b/tests/test_netjson.py index d2a910b..470df56 100644 --- a/tests/test_netjson.py +++ b/tests/test_netjson.py @@ -121,6 +121,10 @@ def test_json_dict(self): self.assertIn('hostname', data['nodes'][0]['properties']) # ensure local_addresses is present self.assertIn('local_addresses', data['nodes'][0]) + # ensure label is present when present in source + self.assertIn('label', data['nodes'][0]) + # ensure label is absent when absent in source + self.assertNotIn('label', data['nodes'][2]) # ensure additional link properties are present self.assertIn('properties', data['links'][0]) self.assertIn('custom_property', data['links'][0]['properties']) diff --git a/tests/test_openvpn.py b/tests/test_openvpn.py index 485dee7..d93ce68 100644 --- a/tests/test_openvpn.py +++ b/tests/test_openvpn.py @@ -1,7 +1,6 @@ import os import networkx -import six from netdiff import OpenvpnParser from netdiff.tests import TestCase