[ONOS-7757] Support onos-local and embedded cluster configurations
- Refactor cluster.json to support internal/external nodes ('controller' and 'storage')
- Bootstrap embedded partitions when 'storage' nodes not present
- Update onos-gen-config script to generate cluster.json based on environment variables
- Update setup scenario to ignore missing $OCC# environment variables

Change-Id: Ia93b64e13d7a7c35ed712da4c681425e3ccf9fe9
diff --git a/core/api/src/main/java/org/onosproject/cluster/ClusterMetadata.java b/core/api/src/main/java/org/onosproject/cluster/ClusterMetadata.java
index 8f7d960..e1f664e 100644
--- a/core/api/src/main/java/org/onosproject/cluster/ClusterMetadata.java
+++ b/core/api/src/main/java/org/onosproject/cluster/ClusterMetadata.java
@@ -44,7 +44,8 @@
     private final ProviderId providerId;
     private final String name;
     private final ControllerNode localNode;
-    private final Set<Node> nodes;
+    private final Set<ControllerNode> controllerNodes;
+    private final Set<Node> storageNodes;
 
     public static final Funnel<ClusterMetadata> HASH_FUNNEL = new Funnel<ClusterMetadata>() {
         @Override
@@ -58,22 +59,26 @@
         providerId = null;
         name = null;
         localNode = null;
-        nodes = null;
+        controllerNodes = null;
+        storageNodes = null;
     }
 
     public ClusterMetadata(
         ProviderId providerId,
         String name,
         ControllerNode localNode,
-        Set<Node> nodes) {
+        Set<ControllerNode> controllerNodes,
+        Set<Node> storageNodes) {
         this.providerId = checkNotNull(providerId);
         this.name = checkNotNull(name);
         this.localNode = localNode;
-        this.nodes = ImmutableSet.copyOf(checkNotNull(nodes));
+        this.controllerNodes = ImmutableSet.copyOf(checkNotNull(controllerNodes));
+        this.storageNodes = ImmutableSet.copyOf(checkNotNull(storageNodes));
     }
 
-    public ClusterMetadata(String name, ControllerNode localNode, Set<Node> nodes) {
-        this(new ProviderId("none", "none"), name, localNode, nodes);
+    public ClusterMetadata(
+            String name, ControllerNode localNode, Set<ControllerNode> controllerNodes, Set<Node> storageNodes) {
+        this(new ProviderId("none", "none"), name, localNode, controllerNodes, storageNodes);
     }
 
     @Override
@@ -102,8 +107,17 @@
      * Returns the collection of {@link org.onosproject.cluster.ControllerNode nodes} that make up the cluster.
      * @return cluster nodes
      */
+    @Deprecated
     public Collection<ControllerNode> getNodes() {
-        return (Collection) nodes;
+        return getControllerNodes();
+    }
+
+    /**
+     * Returns the collection of {@link org.onosproject.cluster.ControllerNode nodes} that make up the cluster.
+     * @return controller nodes
+     */
+    public Collection<ControllerNode> getControllerNodes() {
+        return controllerNodes;
     }
 
     /**
@@ -112,7 +126,7 @@
      * @return the collection of storage nodes
      */
     public Collection<Node> getStorageNodes() {
-        return nodes;
+        return storageNodes;
     }
 
     /**
@@ -131,13 +145,14 @@
         return MoreObjects.toStringHelper(ClusterMetadata.class)
                 .add("providerId", providerId)
                 .add("name", name)
-                .add("nodes", nodes)
+                .add("controllerNodes", controllerNodes)
+                .add("storageNodes", storageNodes)
                 .toString();
     }
 
     @Override
     public int hashCode() {
-        return Arrays.deepHashCode(new Object[] {providerId, name, nodes});
+        return Arrays.deepHashCode(new Object[] {providerId, name, controllerNodes, storageNodes});
     }
 
     /*
@@ -157,8 +172,10 @@
         ClusterMetadata that = (ClusterMetadata) object;
 
         return Objects.equals(this.name, that.name) &&
-               this.localNode.equals(that.localNode) &&
-               Objects.equals(this.nodes.size(), that.nodes.size()) &&
-               Sets.symmetricDifference(this.nodes, that.nodes).isEmpty();
+                this.localNode.equals(that.localNode) &&
+                Objects.equals(this.controllerNodes.size(), that.controllerNodes.size()) &&
+                Sets.symmetricDifference(this.controllerNodes, that.controllerNodes).isEmpty() &&
+                Objects.equals(this.storageNodes.size(), that.storageNodes.size()) &&
+                Sets.symmetricDifference(this.storageNodes, that.storageNodes).isEmpty();
     }
 }
diff --git a/core/api/src/test/java/org/onosproject/cluster/ClusterMetadataEventTest.java b/core/api/src/test/java/org/onosproject/cluster/ClusterMetadataEventTest.java
index b176aee..49ff25f 100644
--- a/core/api/src/test/java/org/onosproject/cluster/ClusterMetadataEventTest.java
+++ b/core/api/src/test/java/org/onosproject/cluster/ClusterMetadataEventTest.java
@@ -38,11 +38,11 @@
     private final ControllerNode n2 =
             new DefaultControllerNode(nid2, IpAddress.valueOf("10.0.0.2"), 9876);
     private final ClusterMetadata metadata1 =
-            new ClusterMetadata("foo", n1, ImmutableSet.of(n1));
+            new ClusterMetadata("foo", n1, ImmutableSet.of(), ImmutableSet.of(n1));
     private final ClusterMetadata metadata2 =
-            new ClusterMetadata("bar", n1, ImmutableSet.of(n1, n2));
+            new ClusterMetadata("bar", n1, ImmutableSet.of(), ImmutableSet.of(n1, n2));
     private final ClusterMetadata metadata3 =
-            new ClusterMetadata("baz", n1, ImmutableSet.of(n2));
+            new ClusterMetadata("baz", n1, ImmutableSet.of(), ImmutableSet.of(n2));
 
     private final ClusterMetadataEvent event1 =
             new ClusterMetadataEvent(ClusterMetadataEvent.Type.METADATA_CHANGED, metadata1, time1);
diff --git a/core/api/src/test/java/org/onosproject/cluster/ClusterMetadataServiceAdapter.java b/core/api/src/test/java/org/onosproject/cluster/ClusterMetadataServiceAdapter.java
index b1b46f9..f4f6100 100644
--- a/core/api/src/test/java/org/onosproject/cluster/ClusterMetadataServiceAdapter.java
+++ b/core/api/src/test/java/org/onosproject/cluster/ClusterMetadataServiceAdapter.java
@@ -27,7 +27,11 @@
     public ClusterMetadata getClusterMetadata() {
         final NodeId nid = new NodeId("test-node");
         final IpAddress addr = IpAddress.valueOf(0);
-        return new ClusterMetadata("test-cluster", new DefaultControllerNode(nid, addr), Sets.newHashSet());
+        return new ClusterMetadata(
+                "test-cluster",
+                new DefaultControllerNode(nid, addr),
+                Sets.newHashSet(),
+                Sets.newHashSet());
     }
 
     @Override
diff --git a/core/api/src/test/java/org/onosproject/cluster/ClusterMetadataTest.java b/core/api/src/test/java/org/onosproject/cluster/ClusterMetadataTest.java
index d6e2f33..1b32363 100644
--- a/core/api/src/test/java/org/onosproject/cluster/ClusterMetadataTest.java
+++ b/core/api/src/test/java/org/onosproject/cluster/ClusterMetadataTest.java
@@ -39,11 +39,11 @@
             new DefaultControllerNode(nid2, IpAddress.valueOf("10.0.0.2"), 9876);
 
     private final ClusterMetadata metadata1 =
-            new ClusterMetadata("foo", n1, ImmutableSet.of(n1));
+            new ClusterMetadata("foo", n1, ImmutableSet.of(), ImmutableSet.of(n1));
     private final ClusterMetadata sameAsMetadata1 =
-            new ClusterMetadata("foo", n1, ImmutableSet.of(n1));
+            new ClusterMetadata("foo", n1, ImmutableSet.of(), ImmutableSet.of(n1));
     private final ClusterMetadata metadata2 =
-            new ClusterMetadata("bar", n1, ImmutableSet.of(n1, n2));
+            new ClusterMetadata("bar", n1, ImmutableSet.of(n1), ImmutableSet.of(n1, n2));
     private final ProviderId defaultProvider =
             new ProviderId("none", "none");
     /**
@@ -64,8 +64,10 @@
     @Test
     public void checkConstruction() {
         assertThat(metadata2.getName(), is("bar"));
-        assertThat(metadata2.getNodes(), hasSize(2));
-        assertThat(metadata2.getNodes(), contains(n1, n2));
+        assertThat(metadata2.getControllerNodes(), hasSize(1));
+        assertThat(metadata2.getControllerNodes(), contains(n1));
+        assertThat(metadata2.getStorageNodes(), hasSize(2));
+        assertThat(metadata2.getStorageNodes(), contains(n1, n2));
         assertThat(metadata1.providerId(), is(defaultProvider));
     }
 }
diff --git a/core/net/src/main/java/org/onosproject/cluster/impl/ClusterMetadataManager.java b/core/net/src/main/java/org/onosproject/cluster/impl/ClusterMetadataManager.java
index 347c812..7dbf139 100644
--- a/core/net/src/main/java/org/onosproject/cluster/impl/ClusterMetadataManager.java
+++ b/core/net/src/main/java/org/onosproject/cluster/impl/ClusterMetadataManager.java
@@ -98,7 +98,8 @@
     public ControllerNode getLocalNode() {
         checkPermission(CLUSTER_READ);
         if (localNode == null) {
-            ControllerNode localNode = getProvider().getClusterMetadata().value().getLocalNode();
+            ClusterMetadata metadata = getProvider().getClusterMetadata().value();
+            ControllerNode localNode = metadata.getLocalNode();
             try {
                 if (localNode != null) {
                     this.localNode = new DefaultControllerNode(
@@ -107,7 +108,15 @@
                         localNode.tcpPort());
                 } else {
                     IpAddress ip = findLocalIp();
-                    this.localNode = new DefaultControllerNode(NodeId.nodeId(ip.toString()), ip);
+                    localNode = metadata.getControllerNodes().stream()
+                        .filter(node -> node.ip().equals(ip))
+                        .findFirst()
+                        .orElse(null);
+                    if (localNode != null) {
+                        this.localNode = localNode;
+                    } else {
+                        this.localNode = new DefaultControllerNode(NodeId.nodeId(ip.toString()), ip);
+                    }
                 }
             } catch (SocketException e) {
                 throw new IllegalStateException(e);
diff --git a/core/net/src/main/java/org/onosproject/cluster/impl/ConfigFileBasedClusterMetadataProvider.java b/core/net/src/main/java/org/onosproject/cluster/impl/ConfigFileBasedClusterMetadataProvider.java
index 8086206..e93ce33 100644
--- a/core/net/src/main/java/org/onosproject/cluster/impl/ConfigFileBasedClusterMetadataProvider.java
+++ b/core/net/src/main/java/org/onosproject/cluster/impl/ConfigFileBasedClusterMetadataProvider.java
@@ -28,6 +28,7 @@
 import java.util.stream.Collectors;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.Sets;
 import com.google.common.io.Files;
 import org.apache.felix.scr.annotations.Activate;
 import org.apache.felix.scr.annotations.Component;
@@ -111,10 +112,14 @@
     private ClusterMetadataPrototype toPrototype(ClusterMetadata metadata) {
         ClusterMetadataPrototype prototype = new ClusterMetadataPrototype();
         prototype.setName(metadata.getName());
-        prototype.setCluster(metadata.getNodes()
-            .stream()
-            .map(this::toPrototype)
-            .collect(Collectors.toSet()));
+        prototype.setController(metadata.getNodes()
+                .stream()
+                .map(this::toPrototype)
+                .collect(Collectors.toSet()));
+        prototype.setStorage(metadata.getStorageNodes()
+                .stream()
+                .map(this::toPrototype)
+                .collect(Collectors.toSet()));
         return prototype;
     }
 
@@ -241,13 +246,20 @@
                         metadata.getNode().getPort() != null
                             ? metadata.getNode().getPort()
                             : DefaultControllerNode.DEFAULT_PORT) : null,
-                metadata.getCluster()
-                    .stream()
-                    .map(node -> new DefaultControllerNode(
-                        NodeId.nodeId(node.getId()),
-                        IpAddress.valueOf(node.getIp()),
-                        node.getPort() != null ? node.getPort() : 5679))
-                    .collect(Collectors.toSet())),
+                    metadata.getController()
+                        .stream()
+                        .map(node -> new DefaultControllerNode(
+                            NodeId.nodeId(node.getId()),
+                            IpAddress.valueOf(node.getIp()),
+                            node.getPort() != null ? node.getPort() : 5679))
+                        .collect(Collectors.toSet()),
+                    metadata.getStorage()
+                        .stream()
+                        .map(node -> new DefaultControllerNode(
+                            NodeId.nodeId(node.getId()),
+                            IpAddress.valueOf(node.getIp()),
+                            node.getPort() != null ? node.getPort() : 5679))
+                        .collect(Collectors.toSet())),
                 version);
         } catch (IOException e) {
             throw new IllegalArgumentException(e);
@@ -278,7 +290,8 @@
     private static class ClusterMetadataPrototype {
         private String name;
         private NodePrototype node;
-        private Set<NodePrototype> cluster;
+        private Set<NodePrototype> controller = Sets.newHashSet();
+        private Set<NodePrototype> storage = Sets.newHashSet();
 
         public String getName() {
             return name;
@@ -296,12 +309,20 @@
             this.node = node;
         }
 
-        public Set<NodePrototype> getCluster() {
-            return cluster;
+        public Set<NodePrototype> getController() {
+            return controller;
         }
 
-        public void setCluster(Set<NodePrototype> cluster) {
-            this.cluster = cluster;
+        public void setController(Set<NodePrototype> controller) {
+            this.controller = controller;
+        }
+
+        public Set<NodePrototype> getStorage() {
+            return storage;
+        }
+
+        public void setStorage(Set<NodePrototype> storage) {
+            this.storage = storage;
         }
     }
 
diff --git a/core/net/src/main/java/org/onosproject/cluster/impl/DefaultClusterMetadataProvider.java b/core/net/src/main/java/org/onosproject/cluster/impl/DefaultClusterMetadataProvider.java
index aa32da2..a6f24ae 100644
--- a/core/net/src/main/java/org/onosproject/cluster/impl/DefaultClusterMetadataProvider.java
+++ b/core/net/src/main/java/org/onosproject/cluster/impl/DefaultClusterMetadataProvider.java
@@ -73,7 +73,7 @@
         ControllerNode localNode =
                 new DefaultControllerNode(new NodeId(localIp), IpAddress.valueOf(localIp), DEFAULT_ONOS_PORT);
         ClusterMetadata metadata = new ClusterMetadata(
-            PROVIDER_ID, "default", localNode, ImmutableSet.of());
+            PROVIDER_ID, "default", localNode, ImmutableSet.of(), ImmutableSet.of());
         long version = System.currentTimeMillis();
         cachedMetadata.set(new Versioned<>(metadata, version));
         providerRegistry.register(this);
diff --git a/core/store/dist/src/test/java/org/onosproject/store/cluster/messaging/impl/NettyMessagingManagerTest.java b/core/store/dist/src/test/java/org/onosproject/store/cluster/messaging/impl/NettyMessagingManagerTest.java
index dfb2f67..ac8a9bc 100644
--- a/core/store/dist/src/test/java/org/onosproject/store/cluster/messaging/impl/NettyMessagingManagerTest.java
+++ b/core/store/dist/src/test/java/org/onosproject/store/cluster/messaging/impl/NettyMessagingManagerTest.java
@@ -245,7 +245,7 @@
             @Override
             public ClusterMetadata getClusterMetadata() {
                 return new ClusterMetadata(new ProviderId(DUMMY_NAME, DUMMY_NAME),
-                                           name, getLocalNode(), Sets.newHashSet());
+                                           name, getLocalNode(), Sets.newHashSet(), Sets.newHashSet());
             }
 
             @Override
diff --git a/core/store/primitives/src/main/java/org/onosproject/store/atomix/impl/AtomixManager.java b/core/store/primitives/src/main/java/org/onosproject/store/atomix/impl/AtomixManager.java
index 24b86f5..ba66372 100644
--- a/core/store/primitives/src/main/java/org/onosproject/store/atomix/impl/AtomixManager.java
+++ b/core/store/primitives/src/main/java/org/onosproject/store/atomix/impl/AtomixManager.java
@@ -15,16 +15,21 @@
  */
 package org.onosproject.store.atomix.impl;
 
+import java.io.File;
+import java.util.Collections;
+import java.util.List;
 import java.util.stream.Collectors;
 
 import io.atomix.cluster.discovery.BootstrapDiscoveryProvider;
 import io.atomix.core.Atomix;
+import io.atomix.protocols.raft.partition.RaftPartitionGroup;
 import org.apache.felix.scr.annotations.Activate;
 import org.apache.felix.scr.annotations.Component;
 import org.apache.felix.scr.annotations.Deactivate;
 import org.apache.felix.scr.annotations.Reference;
 import org.apache.felix.scr.annotations.ReferenceCardinality;
 import org.apache.felix.scr.annotations.Service;
+import org.onosproject.cluster.ClusterMetadata;
 import org.onosproject.cluster.ClusterMetadataService;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -35,7 +40,7 @@
 @Component(immediate = true)
 @Service(value = AtomixManager.class)
 public class AtomixManager {
-
+    private static final String LOCAL_DATA_DIR = System.getProperty("karaf.data") + "/db/partitions/";
     private final Logger log = LoggerFactory.getLogger(getClass());
 
     @Reference(cardinality = ReferenceCardinality.MANDATORY_UNARY)
@@ -54,6 +59,7 @@
 
     @Activate
     public void activate() {
+        log.info("{}", metadataService.getClusterMetadata());
         atomix = createAtomix();
         atomix.start().join();
         log.info("Started");
@@ -66,19 +72,58 @@
     }
 
     private Atomix createAtomix() {
-        return Atomix.builder(getClass().getClassLoader())
-            .withClusterId(metadataService.getClusterMetadata().getName())
-            .withMemberId(metadataService.getLocalNode().id().id())
-            .withAddress(metadataService.getLocalNode().ip().toString(), metadataService.getLocalNode().tcpPort())
-            .withProperty("type", "onos")
-            .withMembershipProvider(BootstrapDiscoveryProvider.builder()
-                .withNodes(metadataService.getClusterMetadata().getStorageNodes().stream()
-                    .map(node -> io.atomix.cluster.Node.builder()
-                        .withId(node.id().id())
-                        .withAddress(node.ip().toString(), node.tcpPort())
-                        .build())
-                    .collect(Collectors.toList()))
-                .build())
-            .build();
+        ClusterMetadata metadata = metadataService.getClusterMetadata();
+        if (!metadata.getStorageNodes().isEmpty()) {
+            // If storage nodes are defined, construct an instance that connects to them for service discovery.
+            return Atomix.builder(getClass().getClassLoader())
+                .withClusterId(metadata.getName())
+                .withMemberId(metadataService.getLocalNode().id().id())
+                .withAddress(metadataService.getLocalNode().ip().toString(), metadataService.getLocalNode().tcpPort())
+                .withProperty("type", "onos")
+                .withMembershipProvider(BootstrapDiscoveryProvider.builder()
+                    .withNodes(metadata.getStorageNodes().stream()
+                        .map(node -> io.atomix.cluster.Node.builder()
+                            .withId(node.id().id())
+                            .withAddress(node.ip().toString(), node.tcpPort())
+                            .build())
+                        .collect(Collectors.toList()))
+                    .build())
+                .build();
+        } else {
+            log.warn("No storage nodes found in cluster metadata!");
+            log.warn("Bootstrapping ONOS cluster in test mode! For production use, configure external storage nodes.");
+
+            // If storage nodes are not defined, construct a local instance with a Raft partition group.
+            List<String> raftMembers = !metadata.getControllerNodes().isEmpty()
+                ? metadata.getControllerNodes()
+                .stream()
+                .map(node -> node.id().id())
+                .collect(Collectors.toList())
+                : Collections.singletonList(metadataService.getLocalNode().id().id());
+            return Atomix.builder(getClass().getClassLoader())
+                .withClusterId(metadata.getName())
+                .withMemberId(metadataService.getLocalNode().id().id())
+                .withAddress(metadataService.getLocalNode().ip().toString(), metadataService.getLocalNode().tcpPort())
+                .withProperty("type", "onos")
+                .withMembershipProvider(BootstrapDiscoveryProvider.builder()
+                    .withNodes(metadata.getControllerNodes().stream()
+                        .map(node -> io.atomix.cluster.Node.builder()
+                            .withId(node.id().id())
+                            .withAddress(node.ip().toString(), node.tcpPort())
+                            .build())
+                        .collect(Collectors.toList()))
+                    .build())
+                .withManagementGroup(RaftPartitionGroup.builder("system")
+                    .withNumPartitions(1)
+                    .withDataDirectory(new File(LOCAL_DATA_DIR, "system"))
+                    .withMembers(raftMembers)
+                    .build())
+                .addPartitionGroup(RaftPartitionGroup.builder("raft")
+                    .withNumPartitions(raftMembers.size())
+                    .withDataDirectory(new File(LOCAL_DATA_DIR, "data"))
+                    .withMembers(raftMembers)
+                    .build())
+                .build();
+        }
     }
 }
diff --git a/tools/package/onos-run-karaf b/tools/package/onos-run-karaf
index d35468a..7c7a8cc 100755
--- a/tools/package/onos-run-karaf
+++ b/tools/package/onos-run-karaf
@@ -55,7 +55,11 @@
     cat > $ONOS_DIR/config/cluster.json <<-EOF
     {
       "name": "default",
-      "cluster": [ {"id": "$IP", "ip": "$IP", "port": 9876 } ]
+      "node": {
+        "id": "$IP",
+        "ip": "$IP",
+        "port": 9876
+      }
     }
 EOF
 
diff --git a/tools/test/bin/onos-gen-config b/tools/test/bin/onos-gen-config
index ef8f2f0..9f31dee 100755
--- a/tools/test/bin/onos-gen-config
+++ b/tools/test/bin/onos-gen-config
@@ -33,68 +33,74 @@
 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_nodes(ips=None, default_port=5679):
-    node = lambda id, ip, port : { 'id': id, 'ip': ip, 'port': port }
+def get_vars(type):
+    types = {'controller': 'OC', 'storage': 'OCC'}
+    vars = []
+    for var in environ:
+        if re.match(r"{}[0-9]+".format(types[type]), var):
+            vars.append(var)
+    return sorted(vars, key=alphanum_key)
+
+
+def get_nodes(type, 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_OCC_vars() ]
+        ips = [environ[v] for v in get_vars(type)]
     i = 1
     for ip_string in ips:
         address_tuple = ip_string.split(":")
         if len(address_tuple) == 3:
-            id=address_tuple[0]
-            ip=address_tuple[1]
-            port=int(address_tuple[2])
+            id = address_tuple[0]
+            ip = address_tuple[1]
+            port = int(address_tuple[2])
         else:
-            id='atomix-{}'.format(i)
+            id = '{}-{}'.format(prefixes[type], i)
             i += 1
-            ip=ip_string
-            port=default_port
+            ip = ip_string
+            port = default_port
         result.append(node(id, ip, port))
     return result
 
+
 if __name__ == '__main__':
-  parser = argparse.ArgumentParser(
-      description="Generate the partitions json file given a list of IPs or from the $OC* environment variables.")
-  parser.add_argument(
-      '-s', '--partition-size', type=int, default=3,
-      help="Number of nodes per partition. Note that partition sizes smaller than 3 are not fault tolerant. Defaults to 3." )
-  parser.add_argument(
-      '-n', '--num-partitions', type=int,
-      help="Number of partitions. Defaults to the number of nodes in the cluster." )
- # TODO: make filename and nodes independent. This will break backwards compatibility with existing usage.
-  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(
-      'nodes', metavar='node_ip', type=str, nargs='*',
-      help='IP Address(es) of the node(s) in the cluster. If no IPs are given, ' +
-           'will use the $OC* environment variables. NOTE: these arguemnts' +
-           ' are only processed after the filename argument.')
+    parser = argparse.ArgumentParser(
+        description="Generate the partitions json file given a list of IPs or from environment variables.")
+    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='+',
+        help='IP Address(es) of the storage nodes. If no IPs are given, ' +
+             'will use the $OCC* environment variables. NOTE: these arguemnts' +
+             ' are only processed after the filename argument.')
 
-  args = parser.parse_args()
-  filename = args.filename
-  partition_size = args.partition_size
-  cluster = get_nodes(args.nodes)
-  num_partitions = args.num_partitions
-  if not num_partitions:
-    num_partitions = len(cluster)
+    args = parser.parse_args()
+    filename = args.filename
+    controller = get_nodes('controller', args.controller_nodes)
+    storage = get_nodes('storage', args.storage_nodes)
 
-  data = {
-           'name': 'onos',
-           'cluster': cluster
-         }
-  output = json.dumps(data, indent=4)
+    if len(storage) > 0:
+        data = {
+            'name': 'onos',
+            'storage': storage
+        }
+    else:
+        data = {
+            'name': 'onos',
+            'controller': controller
+        }
+    output = json.dumps(data, indent=4)
 
-  if filename:
-    with open(filename, 'w') as f:
-      f.write(output)
-  else:
-    print output
+    if filename:
+        with open(filename, 'w') as f:
+            f.write(output)
+    else:
+        print output
diff --git a/tools/test/scenarios/setup.xml b/tools/test/scenarios/setup.xml
index a3bdf32..d148d52 100644
--- a/tools/test/scenarios/setup.xml
+++ b/tools/test/scenarios/setup.xml
@@ -20,26 +20,30 @@
               exec="onos ${OC1} cfg set org.onosproject.net.intent.impl.compiler.IntentConfigurableRegistrator useFlowObjectives true"/>
 
         <group name="Cleanup">
-            <parallel var="${OC#}">
-                <step name="Push-Bits-${#}" exec="onos-push-bits ${OC#}"
-                      unless="${OCT}"/>
-                <step name="Uninstall-${#}" exec="onos-uninstall ${OC#}"/>
-                <step name="Kill-${#}" env="~" exec="onos-kill ${OC#}"
-                      requires="Uninstall-${#}"/>
-            </parallel>
-            <parallel var="${OCC#}">
-                <step name="Atomix-Kill-${#}"
-                      env="~"
-                      exec="atomix-kill ${OCC#}"/>
-                <step name="Atomix-Uninstall-${#}"
-                      exec="atomix-uninstall ${OCC#}"
-                      requires="Atomix-Kill-${#}"/>
-            </parallel>
+            <group name="Cleanup-ONOS">
+                <parallel var="${OC#}">
+                    <step name="Push-Bits-${#}" exec="onos-push-bits ${OC#}"
+                          unless="${OCT}"/>
+                    <step name="Uninstall-${#}" exec="onos-uninstall ${OC#}"/>
+                    <step name="Kill-${#}" env="~" exec="onos-kill ${OC#}"
+                          requires="Uninstall-${#}"/>
+                </parallel>
+            </group>
+            <group name="Cleanup-Atomix" if="${OCC1}">
+                <parallel var="${OCC#}">
+                    <step name="Atomix-Kill-${#}"
+                          env="~"
+                          exec="atomix-kill ${OCC#}"/>
+                    <step name="Atomix-Uninstall-${#}"
+                          exec="atomix-uninstall ${OCC#}"
+                          requires="Atomix-Kill-${#}"/>
+                </parallel>
+            </group>
         </group>
 
-        <group name="Install-Atomix">
-            <step name="Generate-Cluster-Key" exec="onos-gen-cluster-key -f" />
+        <step name="Generate-Cluster-Key" exec="onos-gen-cluster-key -f"/>
 
+        <group name="Install-Atomix" if="${OCC1}">
             <group name="Parallel-Install-Atomix">
                 <parallel var="${OCC#}">
                     <step name="Parallel-Install-Atomix-${#}" exec="atomix-install ${OCC#}"