test: add BGPContainer.define_policy() and BGPContainer.assign_policy()

Signed-off-by: ISHIDA Wataru <ishida.wataru@lab.ntt.co.jp>

3 files changed, 26 insertions, 44 deletions

diff --git a/test/lib/base.py b/test/lib/base.py
index d92f7b9a..d64afdae 100644
--- a/test/lib/base.py
+++ b/test/lib/base.py
@@ -323,21 +323,32 @@ class BGPContainer(Container):
             self.create_config()
             self.reload_config()
 
+    def add_policy(self, policy, peer, typ, default='accept', reload_config=True):
+        self.set_default_policy(peer, typ, default)
+        self.define_policy(policy)
+        self.assign_policy(peer, policy, typ)
+        if self.is_running and reload_config:
+            self.create_config()
+            self.reload_config()
+
     def set_default_policy(self, peer, typ, default):
-        if typ in ['in', 'import', 'export'] and default in ['reject', 'accept']:
+        if typ in ['in', 'out', 'import', 'export'] and default in ['reject', 'accept']:
             if 'default-policy' not in self.peers[peer]:
                 self.peers[peer]['default-policy'] = {}
             self.peers[peer]['default-policy'][typ] = default
         else:
             raise Exception('wrong type or default')
 
-    def add_policy(self, policy, peer=None, reload_config=True):
+    def define_policy(self, policy):
         self.policies[policy['name']] = policy
-        if peer in self.peers:
-            self.peers[peer]['policies'][policy['name']] = policy
-        if self.is_running and reload_config:
-            self.create_config()
-            self.reload_config()
+
+    def assign_policy(self, peer, policy, typ):
+        if peer not in self.peers:
+            raise Exception('peer {0} not found'.format(peer.name))
+        name = policy['name']
+        if name not in self.policies:
+            raise Exception('policy {0} not found'.format(name))
+        self.peers[peer]['policies'][typ] = policy
 
     def get_local_rib(self, peer, rf):
         raise Exception('implement get_local_rib() method')
diff --git a/test/lib/gobgp.py b/test/lib/gobgp.py
index 0ae2d71b..5b7a847e 100644
--- a/test/lib/gobgp.py
+++ b/test/lib/gobgp.py
@@ -281,27 +281,11 @@ class GoBGPContainer(BGPContainer):
                 n['route-reflector'] = {'config' : {'route-reflector-client': True,
                                                    'route-reflector-cluster-id': clusterId}}
 
-            f = lambda typ: [p for p in info['policies'].itervalues() if p['type'] == typ]
-            import_policies = f('import')
-            export_policies = f('export')
-            in_policies = f('in')
-            f = lambda typ: [p['default'] for p in info['policies'].itervalues() if p['type'] == typ and 'default' in p]
-            default_import_policy = f('import')
-            default_export_policy = f('export')
-            default_in_policy  = f('in')
-
-            if len(import_policies) + len(export_policies) + len(in_policies) + len(default_import_policy) \
-                + len(default_export_policy) + len(default_in_policy) > 0:
+            if len(info.get('default-policy', [])) + len(info.get('policies', [])) > 0:
                 n['apply-policy'] = {'config': {}}
 
-            if len(import_policies) > 0:
-                n['apply-policy']['config']['import-policy-list'] = [p['name'] for p in import_policies]
-
-            if len(export_policies) > 0:
-                n['apply-policy']['config']['export-policy-list'] = [p['name'] for p in export_policies]
-
-            if len(in_policies) > 0:
-                n['apply-policy']['config']['in-policy-list'] = [p['name'] for p in in_policies]
+            for typ, p in info.get('policies', {}).iteritems():
+                n['apply-policy']['config']['{0}-policy-list'.format(typ)] = [p['name']]
 
             def f(v):
                 if v == 'reject':
@@ -310,20 +294,8 @@ class GoBGPContainer(BGPContainer):
                     return 'accept-route'
                 raise Exception('invalid default policy type {0}'.format(v))
 
-            if len(default_import_policy) > 0:
-               n['apply-policy']['config']['default-import-policy'] = f(default_import_policy[0])
-
-            if len(default_export_policy) > 0:
-               n['apply-policy']['config']['default-export-policy'] = f(default_export_policy[0])
-
-            if len(default_in_policy) > 0:
-               n['apply-policy']['config']['default-in-policy'] = f(default_in_policy[0])
-
-            for typ in ['in', 'import', 'export']:
-                if 'default-policy' in info and typ in info['default-policy']:
-                    if 'apply-policy' not in n:
-                        n['apply-policy'] = {'config': {}}
-                    n['apply-policy']['config']['default-{0}-policy'.format(typ)] = f(info['default-policy'][typ])
+            for typ, d in info.get('default-policy', {}).iteritems():
+                n['apply-policy']['config']['default-{0}-policy'.format(typ)] = f(d)
 
             if 'neighbors' not in config:
                 config['neighbors'] = []
diff --git a/test/lib/quagga.py b/test/lib/quagga.py
index 59cf58e9..ce53a29b 100644
--- a/test/lib/quagga.py
+++ b/test/lib/quagga.py
@@ -164,10 +164,9 @@ class QuaggaBGPContainer(BGPContainer):
             c << 'neighbor {0} remote-as {1}'.format(n_addr, peer.asn)
             if info['is_rs_client']:
                 c << 'neighbor {0} route-server-client'.format(n_addr)
-            for name, policy in info['policies'].iteritems():
-                direction = policy['direction']
-                c << 'neighbor {0} route-map {1} {2}'.format(n_addr, name,
-                                                             direction)
+            for typ, p in info['policies'].iteritems():
+                c << 'neighbor {0} route-map {1} {2}'.format(n_addr, p['name'],
+                                                             typ)
             if info['passwd']:
                 c << 'neighbor {0} password {1}'.format(n_addr, info['passwd'])
             if info['passive']: