BgpRoute: add support for sourceProtocol (#356)

batfish · Apr 24, 2019 · 8430a0f · 8430a0f
1 parent 4b8380b
commit 8430a0f
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 15 deletions.
diff --git a/pybatfish/datamodel/route.py b/pybatfish/datamodel/route.py
@@ -36,6 +36,7 @@ class BgpRoute(DataModelElement):
     :ivar metric: The metric of the route.
     :ivar originatorIp: The IP address of the originator of the route.
     :ivar originType: The origin type of the route.
+    :ivar sourceProtocol: The source protocol of the route.
     """
 
     network = attr.ib(type=str)
@@ -46,6 +47,7 @@ class BgpRoute(DataModelElement):
     communities = attr.ib(type=list, default=[])
     localPreference = attr.ib(type=int, default=0)
     metric = attr.ib(type=int, default=0)
+    sourceProtocol = attr.ib(type=str, default=None)
 
     @classmethod
     def from_dict(cls, json_dict):
@@ -57,7 +59,8 @@ def from_dict(cls, json_dict):
                         json_dict.get("asPath", []),
                         json_dict.get("communities", []),
                         json_dict.get("localPreference", 0),
-                        json_dict.get("metric", 0))
+                        json_dict.get("metric", 0),
+                        json_dict.get("srcProtocol", None))
 
     def dict(self):
         # type: () -> Dict
@@ -71,7 +74,8 @@ def dict(self):
             'metric': self.metric,
             'originatorIp': self.originatorIp,
             'originType': self.originType,
-            'protocol': self.protocol
+            'protocol': self.protocol,
+            'srcProtocol': self.sourceProtocol
         }
 
     def _repr_html_(self):
@@ -89,6 +93,7 @@ def _repr_html_lines(self):
         lines.append('Originator IP: %s' % self.originatorIp)
         lines.append('Origin Type: %s' % self.originType)
         lines.append('Protocol: %s' % self.protocol)
+        lines.append('Source Protocol: %s' % self.sourceProtocol)
         return lines
 
 

diff --git a/tests/datamodel/test_route.py b/tests/datamodel/test_route.py
@@ -59,27 +59,30 @@ def testBgpRouteSerialization():
     assert dct['protocol'] == protocol
 
 
-def testBgpRouteStr():
-    bgpRoute = BgpRoute(
+def test_bgp_route_str():
+    bgp_route = BgpRoute(
         network="A",
         asPath=[1, 2],
         communities=[1, 2, 3],
         localPreference=4,
         metric=5,
         originatorIp='1.1.1.1',
         originType='egp',
-        protocol='bgp'
+        protocol='bgp',
+        sourceProtocol='connected'
     )
-    lines = bgpRoute._repr_html_lines()
-    assert len(lines) == 8
-    assert lines[0] == "Network: A"
-    assert lines[1] == "AS Path: [1, 2]"
-    assert lines[2] == "Communities: [1, 2, 3]"
-    assert lines[3] == "Local Preference: 4"
-    assert lines[4] == "Metric: 5"
-    assert lines[5] == "Originator IP: 1.1.1.1"
-    assert lines[6] == "Origin Type: egp"
-    assert lines[7] == "Protocol: bgp"
+    lines = bgp_route._repr_html_lines()
+    assert lines == [
+        "Network: A",
+        "AS Path: [1, 2]",
+        "Communities: [1, 2, 3]",
+        "Local Preference: 4",
+        "Metric: 5",
+        "Originator IP: 1.1.1.1",
+        "Origin Type: egp",
+        "Protocol: bgp",
+        "Source Protocol: connected",
+    ]
 
 
 def testBgpRouteDiffDeserialization():