diff --git a/tools/test/bin/atomix-gen-config b/tools/test/bin/atomix-gen-config
index 5e42dfc..e69f441 100755
--- a/tools/test/bin/atomix-gen-config
+++ b/tools/test/bin/atomix-gen-config
@@ -33,23 +33,32 @@
 convert = lambda text: int(text) if text.isdigit() else text.lower()
 alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]
 
-def get_OCC_vars():
-  vars = []
-  for var in environ:
-    if re.match(r"OCC[0-9]+", var):
-      vars.append(var)
-  return sorted(vars, key=alphanum_key)
+def get_vars_by_type(type):
+    vars = []
+    for var in environ:
+        if re.match(r"{}[0-9]+".format(type), var):
+            vars.append(var)
+    return sorted(vars, key=alphanum_key)
+
+
+def get_vars():
+    vars = get_vars_by_type('OCC')
+    if len(vars) == 0:
+        vars = get_vars_by_type('OC')
+    return vars
+
 
 def get_local_node(node, ips=None):
     if not ips:
-        ips = [ environ[v] for v in get_OCC_vars() ]
+        ips = [environ[v] for v in get_vars()]
     return 'atomix-{}'.format(ips.index(node) + 1)
 
+
 def get_nodes(ips=None, default_port=5679):
     node = lambda id, ip, port: {'id': id, 'address': '{}:{}'.format(ip, port)}
     result = []
     if not ips:
-        ips = [ environ[v] for v in get_OCC_vars() ]
+        ips = [environ[v] for v in get_vars()]
     i = 1
     for ip_string in ips:
         address_tuple = ip_string.split(":")
@@ -65,11 +74,10 @@
         result.append(node(id, ip, port))
     return result
 
+
 def get_local_address(node, ips=None, default_port=5679):
-    result = []
     if not ips:
-        ips = [ environ[v] for v in get_OCC_vars() ]
-    i = 1
+        ips = [environ[v] for v in get_vars()]
     for ip_string in ips:
         address_tuple = ip_string.split(":")
         if len(address_tuple) == 3:
diff --git a/tools/test/bin/onos-config b/tools/test/bin/onos-config
index 772df3a..65df721 100755
--- a/tools/test/bin/onos-config
+++ b/tools/test/bin/onos-config
@@ -66,7 +66,7 @@
 
 # Generate a default cluster.json from the ON* environment variables
 CDEF_FILE=/tmp/${remote}.cluster.json
-onos-gen-config $CDEF_FILE
+onos-gen-config $node $CDEF_FILE
 scp -q $CDEF_FILE $remote_scp:$ONOS_INSTALL_DIR/config/cluster.json
 
 # Copy tools/package/config/ to remote
diff --git a/tools/test/bin/onos-gen-config b/tools/test/bin/onos-gen-config
index 9f31dee..b73df39 100755
--- a/tools/test/bin/onos-gen-config
+++ b/tools/test/bin/onos-gen-config
@@ -34,21 +34,26 @@
 alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]
 
 
-def get_vars(type):
-    types = {'controller': 'OC', 'storage': 'OCC'}
+def get_vars_by_type(type):
     vars = []
     for var in environ:
-        if re.match(r"{}[0-9]+".format(types[type]), var):
+        if re.match(r"{}[0-9]+".format(type), var):
             vars.append(var)
     return sorted(vars, key=alphanum_key)
 
 
-def get_nodes(type, ips=None, default_port=5679):
+def get_vars():
+    vars = get_vars_by_type('OCC')
+    if len(vars) == 0:
+        vars = get_vars_by_type('OC')
+    return vars
+
+
+def get_nodes(ips=None, default_port=5679):
     node = lambda id, ip, port: {'id': id, 'ip': ip, 'port': port}
-    prefixes = {'controller': 'onos', 'storage': 'atomix'}
     result = []
     if not ips:
-        ips = [environ[v] for v in get_vars(type)]
+        ips = [environ[v] for v in get_vars()]
     i = 1
     for ip_string in ips:
         address_tuple = ip_string.split(":")
@@ -57,7 +62,7 @@
             ip = address_tuple[1]
             port = int(address_tuple[2])
         else:
-            id = '{}-{}'.format(prefixes[type], i)
+            id = 'atomix-{}'.format(i)
             i += 1
             ip = ip_string
             port = default_port
@@ -69,34 +74,32 @@
     parser = argparse.ArgumentParser(
         description="Generate the partitions json file given a list of IPs or from environment variables.")
     parser.add_argument(
+        'node', metavar='node_ip', type=str, nargs='?',
+        help='The node for which to generate the configuration')
+    parser.add_argument(
         'filename', metavar='filename', type=str, nargs='?',
         help='File to write output to. If none is provided, output is written to stdout.')
     parser.add_argument(
-        '--controller-nodes', '-c', metavar='node_ip', type=str, nargs='+',
-        help='IP Address(es) of the controller nodes. If no IPs are given, ' +
-             'will use the $OC* environment variables. NOTE: these arguemnts' +
-             ' are only processed after the filename argument.')
-    parser.add_argument(
-        '--storage-nodes', '-s', metavar='node_ip', type=str, nargs='+',
+        '--nodes', '-n', metavar='node_ip', type=str, nargs='+',
         help='IP Address(es) of the storage nodes. If no IPs are given, ' +
-             'will use the $OCC* environment variables. NOTE: these arguemnts' +
+             'will use the $OCC* or $OC* environment variables. NOTE: these arguments' +
              ' are only processed after the filename argument.')
 
     args = parser.parse_args()
+    node = args.node
+    print node
     filename = args.filename
-    controller = get_nodes('controller', args.controller_nodes)
-    storage = get_nodes('storage', args.storage_nodes)
+    nodes = get_nodes(args.nodes)
 
-    if len(storage) > 0:
-        data = {
-            'name': 'onos',
-            'storage': storage
-        }
-    else:
-        data = {
-            'name': 'onos',
-            'controller': controller
-        }
+    data = {
+        'name': 'onos',
+        'node': {
+            'id': node,
+            'ip': node,
+            'port': 9876
+        },
+        'storage': nodes
+    }
     output = json.dumps(data, indent=4)
 
     if filename: