From 52a290f2cc00d19d0ccbfe1314c0a090cff1695c Mon Sep 17 00:00:00 2001
From: Zack Cerza <zack@redhat.com>
Date: Mon, 17 Jan 2022 13:38:10 -0700
Subject: [PATCH] Cluster: Add filter() function

Signed-off-by: Zack Cerza <zack@redhat.com>
---
 teuthology/orchestra/cluster.py           | 13 +++++++++++++
 teuthology/orchestra/test/test_cluster.py | 13 +++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/teuthology/orchestra/cluster.py b/teuthology/orchestra/cluster.py
index dbc74de021..450f469dd9 100644
--- 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 46e3d4c7b4..abbfc6ab2d 100644
--- 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 """
-- 
2.39.5