Update trellisp4.py to use Stratum-enabled bmv2 switch class from ONOS

Change-Id: I6812738d74c430200c7c3d2b31606d2668eb9174
diff --git a/trellis/trellisp4.json b/trellis/trellisp4.json
index 608e033..a12624b 100644
--- a/trellis/trellisp4.json
+++ b/trellis/trellisp4.json
@@ -1,92 +1,104 @@
 {
-    "ports" : {
-        "device:bmv2:s204/3" : {
-            "interfaces" : [
-                {
-                    "name": "h1",
-                    "ips" : [ "10.0.2.254/24" ],
-                    "vlan-untagged": 10
-                }
-            ]
-        },
-        "device:bmv2:s204/4" : {
-            "interfaces" : [
-                {
-                    "name": "h2",
-                    "ips" : [ "10.0.2.254/24" ],
-                    "vlan-tagged": [10]
-                }
-            ]
-        },
-        "device:bmv2:s205/3" : {
-            "interfaces" : [
-                {
-                    "name": "h3",
-                    "ips" : [ "10.0.3.254/24" ],
-                    "vlan-untagged": 20
-                }
-            ]
-        },
-        "device:bmv2:s205/4" : {
-            "interfaces" : [
-                {
-                    "name": "h4",
-                    "ips" : [ "10.0.3.254/24" ],
-                    "vlan-tagged": [20]
-                }
-            ]
+  "ports": {
+    "device:s204/3": {
+      "interfaces": [
+        {
+          "name": "h1",
+          "ips": [
+            "10.0.2.254/24"
+          ],
+          "vlan-untagged": 10
         }
+      ]
     },
-    "devices" : {
-        "device:bmv2:s204" : {
-            "segmentrouting" : {
-                "name" : "s204",
-                "ipv4NodeSid" : 204,
-                "ipv4Loopback" : "192.168.0.204",
-                "ipv6NodeSid" : 214,
-                "ipv6Loopback" : "2000::c0a8:0204",
-                "routerMac" : "00:00:00:00:02:04",
-                "isEdgeRouter" : true,
-                "adjacencySids" : []
-            }
-        },
-        "device:bmv2:s205" : {
-            "segmentrouting" : {
-                "name" : "s205",
-                "ipv4NodeSid" : 205,
-                "ipv4Loopback" : "192.168.0.205",
-                "ipv6NodeSid" : 215,
-                "ipv6Loopback" : "2000::c0a8:0205",
-                "routerMac" : "00:00:00:00:02:05",
-                "isEdgeRouter" : true,
-                "adjacencySids" : []
-            }
-        },
-        "device:bmv2:s226" : {
-            "segmentrouting" : {
-                "name" : "s226",
-                "ipv4NodeSid" : 226,
-                "ipv4Loopback" : "192.168.0.226",
-                "ipv6NodeSid" : 236,
-                "ipv6Loopback" : "2000::c0a8:0226",
-                "routerMac" : "00:00:00:00:02:26",
-                "isEdgeRouter" : false,
-                "adjacencySids" : []
-            }
-        },
-        "device:bmv2:s227" : {
-            "segmentrouting" : {
-                "name" : "s227",
-                "ipv4NodeSid" : 227,
-                "ipv4Loopback" : "192.168.0.227",
-                "ipv6NodeSid" : 237,
-                "ipv6Loopback" : "2000::c0a8:0227",
-                "routerMac" : "00:00:00:00:02:27",
-                "isEdgeRouter" : false,
-                "adjacencySids" : []
-            }
+    "device:s204/4": {
+      "interfaces": [
+        {
+          "name": "h2",
+          "ips": [
+            "10.0.2.254/24"
+          ],
+          "vlan-tagged": [
+            10
+          ]
         }
+      ]
     },
-    "apps" : {
+    "device:s205/3": {
+      "interfaces": [
+        {
+          "name": "h3",
+          "ips": [
+            "10.0.3.254/24"
+          ],
+          "vlan-untagged": 20
+        }
+      ]
+    },
+    "device:s205/4": {
+      "interfaces": [
+        {
+          "name": "h4",
+          "ips": [
+            "10.0.3.254/24"
+          ],
+          "vlan-tagged": [
+            20
+          ]
+        }
+      ]
     }
+  },
+  "devices": {
+    "device:s204": {
+      "segmentrouting": {
+        "name": "s204",
+        "ipv4NodeSid": 204,
+        "ipv4Loopback": "192.168.0.204",
+        "ipv6NodeSid": 214,
+        "ipv6Loopback": "2000::c0a8:0204",
+        "routerMac": "00:00:00:00:02:04",
+        "isEdgeRouter": true,
+        "adjacencySids": []
+      }
+    },
+    "device:s205": {
+      "segmentrouting": {
+        "name": "s205",
+        "ipv4NodeSid": 205,
+        "ipv4Loopback": "192.168.0.205",
+        "ipv6NodeSid": 215,
+        "ipv6Loopback": "2000::c0a8:0205",
+        "routerMac": "00:00:00:00:02:05",
+        "isEdgeRouter": true,
+        "adjacencySids": []
+      }
+    },
+    "device:s226": {
+      "segmentrouting": {
+        "name": "s226",
+        "ipv4NodeSid": 226,
+        "ipv4Loopback": "192.168.0.226",
+        "ipv6NodeSid": 236,
+        "ipv6Loopback": "2000::c0a8:0226",
+        "routerMac": "00:00:00:00:02:26",
+        "isEdgeRouter": false,
+        "adjacencySids": []
+      }
+    },
+    "device:s227": {
+      "segmentrouting": {
+        "name": "s227",
+        "ipv4NodeSid": 227,
+        "ipv4Loopback": "192.168.0.227",
+        "ipv6NodeSid": 237,
+        "ipv6Loopback": "2000::c0a8:0227",
+        "routerMac": "00:00:00:00:02:27",
+        "isEdgeRouter": false,
+        "adjacencySids": []
+      }
+    }
+  },
+  "apps": {
+  }
 }
diff --git a/trellis/trellisp4.py b/trellis/trellisp4.py
index 1adae81..d569976 100755
--- a/trellis/trellisp4.py
+++ b/trellis/trellisp4.py
@@ -1,34 +1,44 @@
 #!/usr/bin/python
 
+import argparse
 import os
 import sys
-import argparse
+
 sys.path.append('..')
 
 from mininet.topo import Topo
 from mininet.net import Mininet
 from mininet.cli import CLI
 from mininet.log import setLogLevel
-from mininet.node import Host, RemoteController
+from mininet.node import RemoteController
 from routinglib import RoutedHost
 from trellislib import TaggedRoutedHost
-from bmv2 import ONOSBmv2Switch
+
+try:
+    from bmv2 import ONOSBmv2Switch, ONOSStratumSwitch
+except ImportError as e:
+    if os.getenv("ONOS_ROOT"):
+        sys.path.append("%s/tools/dev/mininet" % os.getenv("ONOS_ROOT"))
+        from bmv2 import ONOSBmv2Switch, ONOSStratumSwitch
+    else:
+        raise e
 
 PIPECONF_ID = 'org.onosproject.pipelines.fabric'
 
-class Trellis( Topo ):
+
+class Trellis(Topo):
     "Trellis basic topology"
 
-    def __init__( self, *args, **kwargs ):
-        Topo.__init__( self, *args, **kwargs )
+    def __init__(self, *args, **kwargs):
+        Topo.__init__(self, *args, **kwargs)
 
         # Spines
-        s226 = self.addSwitch('s226', cls=ONOSBmv2Switch, grpcport=55226, pipeconf=PIPECONF_ID, portcfg=True)
-        s227 = self.addSwitch('s227', cls=ONOSBmv2Switch, grpcport=55227, pipeconf=PIPECONF_ID, portcfg=True)
+        s226 = self.addP4Switch('s226')
+        s227 = self.addP4Switch('s227')
 
         # Leaves
-        s204 = self.addSwitch('s204', cls=ONOSBmv2Switch, grpcport=55204, pipeconf=PIPECONF_ID, portcfg=True)
-        s205 = self.addSwitch('s205', cls=ONOSBmv2Switch, grpcport=55205, pipeconf=PIPECONF_ID, portcfg=True)
+        s204 = self.addP4Switch('s204')
+        s205 = self.addP4Switch('s205')
 
         # Switch Links
         self.addLink(s226, s204)
@@ -39,34 +49,54 @@
         # NOTE avoid using 10.0.1.0/24 which is the default subnet of quaggas
         # NOTE avoid using 00:00:00:00:00:xx which is the default mac of host behind upstream router
         # IPv4 Hosts
-        h1 = self.addHost('h1', cls=RoutedHost, mac='00:aa:00:00:00:01', ips=['10.0.2.1/24'], gateway='10.0.2.254')
-        h2 = self.addHost('h2', cls=TaggedRoutedHost, mac='00:aa:00:00:00:02', ips=['10.0.2.2/24'], gateway='10.0.2.254', vlan=10)
-        h3 = self.addHost('h3', cls=RoutedHost, mac='00:aa:00:00:00:03', ips=['10.0.3.1/24'], gateway='10.0.3.254')
-        h4 = self.addHost('h4', cls=TaggedRoutedHost, mac='00:aa:00:00:00:04', ips=['10.0.3.2/24'], gateway='10.0.3.254', vlan=20)
+        h1 = self.addHost('h1', cls=RoutedHost, mac='00:aa:00:00:00:01',
+                          ips=['10.0.2.1/24'], gateway='10.0.2.254')
+        h2 = self.addHost('h2', cls=TaggedRoutedHost, mac='00:aa:00:00:00:02',
+                          ips=['10.0.2.2/24'], gateway='10.0.2.254', vlan=10)
+        h3 = self.addHost('h3', cls=RoutedHost, mac='00:aa:00:00:00:03',
+                          ips=['10.0.3.1/24'], gateway='10.0.3.254')
+        h4 = self.addHost('h4', cls=TaggedRoutedHost, mac='00:aa:00:00:00:04',
+                          ips=['10.0.3.2/24'], gateway='10.0.3.254', vlan=20)
         self.addLink(h1, s204)
         self.addLink(h2, s204)
         self.addLink(h3, s205)
         self.addLink(h4, s205)
 
-topos = { 'trellis' : Trellis }
+    def addP4Switch(self, name):
+        return self.addSwitch(name=name,
+                              pipeconf=PIPECONF_ID,
+                              portcfg=True,
+                              onosdevid="device:" + name)
+
+
+topos = {'trellis': Trellis}
+
 
 def main(args):
     topo = Trellis()
     controller = RemoteController('c0', ip=args.onos_ip)
 
-    net = Mininet(topo=topo, controller=None)
+    if args.agent == "stratum":
+        switch = ONOSStratumSwitch
+    else:
+        switch = ONOSBmv2Switch
+
+    net = Mininet(topo=topo, switch=switch, controller=None)
     net.addController(controller)
 
     net.start()
     CLI(net)
     net.stop()
 
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description='BMv2 mininet demo script (2 by 2 fabric)')
     parser.add_argument('--onos-ip', help='ONOS-BMv2 controller IP address',
                         type=str, action="store", required=True)
+    parser.add_argument("-a", "--agent",
+                        help="Agent to use on Bmv2 (pi or stratum)",
+                        required=False, default="pi")
     args = parser.parse_args()
-    setLogLevel('debug')
-
+    setLogLevel('info')
     main(args)