UPS refactoring in progress

Author: belavina

enginecore/enginecore/model/graph_reference.py

@@ -5,7 +5,7 @@
 from enginecore.state.utils import format_as_redis_key
 
 class GraphReference():
-
+    """Graph DB wrapper """
     def __init__(self):
         self._driver = GraphDatabase.driver(
             'bolt://localhost', 
@@ -117,7 +117,8 @@ def get_asset_oid_by_name(cls, session, asset_key, oid_name):
             asset_key(int): key of the asset OID belongs to
             oid_name(str): OID name
         Returns:
-            tuple: str as SNMP OID that belongs to the asset, followed by an int as datatype, followed by optional state details; 
+            tuple: str as SNMP OID that belongs to the asset, 
+                   followed by an int as datatype, followed by optional state details; 
                    returns None if there's no such OID
         """
 
@@ -135,10 +136,12 @@ def get_asset_oid_by_name(cls, session, asset_key, oid_name):
         details = record.get('oid') if record else None
 
         oid_info = details['OID'] if details else None
-        vendor_specs = {v:k for k, v in dict(record['oid_details']).items()} if (record and record['oid_details']) else None
+
+        # get vendor specific information
+        v_specs = {v:k for k, v in dict(record['oid_details']).items()} if (record and record['oid_details']) else None
         oid_data_type = details['dataType'] if oid_info else None
 
-        return oid_info, oid_data_type, vendor_specs
+        return oid_info, oid_data_type, v_specs
 
     @classmethod
     def get_component_oid_by_name(cls, session, component_key, oid_name):
@@ -154,10 +157,10 @@ def get_component_oid_by_name(cls, session, component_key, oid_name):
 
         result = session.run(
             """
-            MATCH (component:Component { key: $key})<-[:HAS_COMPONENT]-(p:Asset)-[:HAS_OID]->(oid:OID {OIDName: $oid_name})
+            MATCH (component:Component { key: $key})<-[:HAS_COMPONENT]-(p:Asset)-[:HAS_OID]->(oid:OID {OIDName: $name})
             RETURN oid, p.key as parent_key
             """,
-            oid_name=oid_name, key=component_key
+            name=oid_name, key=component_key
         )
         record = result.single()
 
@@ -233,8 +236,8 @@ def get_assets_and_connections(cls, session, flatten=True):
                 )
 
                 asset['parent'] = []
-                for r in presults:
-                    asset['parent'].append(dict(r['parent']))
+                for presult in presults:
+                    asset['parent'].append(dict(presult['parent']))
 
             ## Set asset children 
             # format asset children as list of child_key: { child_info }

enginecore/enginecore/model/system_modeler.py

@@ -31,6 +31,11 @@ def _get_set_stm(attr, node_name="asset", supported_attr=[]):
     return ','.join(map(lambda k: "{}.{}={}".format(node_name, _to_camelcase(k), repr(existing[k])), existing))
 
 
+def _get_oid_desc_stm(oid_desc, oid_name):
+    """Format dict attributes as neo4j props"""
+    return 'OIDName: "{}",'.format(oid_name) + ','.join(map(lambda k: '{}: "{}"'.format(oid_desc[k], k), oid_desc))
+
+
 def _add_psu(key, psu_index, attr):
     """Add a PSU to an existing server
     
@@ -277,45 +282,27 @@ def create_ups(key, attr, preset_file=os.path.join(os.path.dirname(__file__), 'p
     preset_file = attr['snmp_preset'] if 'snmp_preset' in attr and attr['snmp_preset'] else preset_file
 
     with open(preset_file) as f, graph_ref.get_session() as session:
+        
+        query = []
         data = json.load(f)
 
-        name = attr['name'] if 'name' in attr and attr['name'] else data['assetName']
-        
-        session.run("\
-        CREATE (:Asset:UPS:SNMPSim { \
-            name: $name,\
-            key: $key,\
-            staticOidFile: $oid_file,\
-            outputPowerCapacity: $pc,\
-            minPowerOnBatteryLevel: $minbat,\
-            fullRechargeTime: $rechargehrs, \
-            type: 'ups',\
-            runtime: $runtime,\
-            port: $port \
-        })", 
-        key=key, 
-        name=name,
-        oid_file=data['staticOidFile'],
-        pc=data['outputPowerCapacity'],
-        minbat=data['minPowerOnBatteryLevel'],
-        rechargehrs=data['fullRechargeTime'],
-        runtime=json.dumps(data['modelRuntime'], sort_keys=True),
-        port=attr['port']
-        )
-
-        set_properties(key, attr)
+        attr['name'] = attr['name'] if 'name' in attr and attr['name'] else data['assetName']
+        attr['runtime'] = json.dumps(data['modelRuntime'], sort_keys=True)
+
+        s_attr = ["name", "type", "key", "off_delay", "on_delay", "staticOidFile", "port", "host",
+                  "outputPowerCapacity", "minPowerOnBatteryLevel", "fullRechargeTime", "runtime"]
+
+        props_stm = _get_props_stm({**attr, **data, **{'key': key, 'type': 'ups'}}, supported_attr=s_attr)
+
+        # create UPS asset
+        query.append("CREATE (ups:Asset:UPS:SNMPSim {{ {} }})".format(props_stm))
 
-        # add batteries
-        session.run("\
-        MATCH (ups:UPS {key: $key})\
-        CREATE (bat:UPSBattery:Battery { \
-            name: $name,\
-            type: 'battery'\
-        })\
-        CREATE (ups)-[:HAS_BATTERY]->(bat)\
-        CREATE (ups)-[:POWERED_BY]->(bat)\
-        ",key=key, name='bat1'
-        )
+        # add batteries (only one for now)
+        query.append("CREATE (bat:UPSBattery:Battery { name: 'bat1', type: 'battery' })")
+        query.append("CREATE (ups)-[:HAS_BATTERY]->(bat)")
+        query.append("CREATE (ups)-[:POWERED_BY]->(bat)")
+
+        session.run("\n".join(query))
 
         for k, v in data["OIDs"].items():
             if k == 'SerialNumber':
@@ -490,15 +477,6 @@ def create_ups(key, attr, preset_file=os.path.join(os.path.dirname(__file__), 'p
             outkey=int("{}{}".format(key,str(i+1))))
 
 
-def _get_oid_desc_stm(oid_desc, oid_name):
-    """Format dict attributes as neo4j props"""
-
-    # existing = dict(
-    #     filter(lambda k: attr[k[0]] != None and (not supported_attr or k[0] in supported_attr), oid_desc.items())
-    # )
-    return 'OIDName: "{}",'.format(oid_name) + ','.join(map(lambda k: '{}: "{}"'.format(oid_desc[k], k), oid_desc))
-
-
 def create_pdu(key, attr, preset_file=os.path.join(os.path.dirname(__file__), 'presets/apc_pdu.json')):
     """Add PDU to the model """