Cluster: Add filter() function

Signed-off-by: Zack Cerza <zack@redhat.com>

diff --git a/teuthology/orchestra/cluster.py b/teuthology/orchestra/cluster.py
index dbc74de021dbbaa98b18dce3f023772426d41db4..450f469dd97ff4c515db736ebeb60a5b2df99ac9 100644 (file)
--- a/teuthology/orchestra/cluster.py
+++ b/teuthology/orchestra/cluster.py
@@ -144,3 +144,16 @@ class Cluster(object):
             if remote not in matches.remotes:
                 c.add(remote, has_roles)
         return c
+
+    def filter(self, func):
+        """
+        Return a cluster whose remotes are filtered by `func`.
+
+        Example::
+            cluster = ctx.cluster.filter(lambda r: r.is_online)
+        """
+        result = self.__class__()
+        for rem, roles in self.remotes.items():
+            if func(rem):
+                result.add(rem, roles)
+        return result

diff --git a/teuthology/orchestra/test/test_cluster.py b/teuthology/orchestra/test/test_cluster.py
index 46e3d4c7b4965acec5de99d71c581c93f5bd16d6..abbfc6ab2da9e2dac5d44848003bd0a6061790c4 100644 (file)
--- a/teuthology/orchestra/test/test_cluster.py
+++ b/teuthology/orchestra/test/test_cluster.py
@@ -187,6 +187,19 @@ class TestCluster(object):
         c_foo = c.exclude('foo', lambda role: role.startswith('b'))
         assert c_foo.remotes == {r2: ['bar'], r3: ['foo']}
 
+    def test_filter(self):
+        r1 = Mock(_name='r1')
+        r2 = Mock(_name='r2')
+        def func(r):
+            return r._name == "r1"
+        c = cluster.Cluster(remotes=[
+                (r1, ['foo']),
+                (r2, ['bar']),
+            ])
+        assert c.filter(func).remotes == {
+            r1: ['foo']
+        }
+
 
 class TestWriteFile(object):
     """ Tests for cluster.write_file """