diff --git a/.ci/tests/examples/is_success.py b/.ci/tests/examples/is_success.py
deleted file mode 100644
index 9493a7377..000000000
--- a/.ci/tests/examples/is_success.py
+++ /dev/null
@@ -1,32 +0,0 @@
-import sys
-from time import sleep
-
-import pymongo
-
-N_ROUNDS=3
-RETRIES=18
-SLEEP=10
-
-def _eprint(*args, **kwargs):
-    print(*args, file=sys.stderr, **kwargs)
-
-def _wait_n_rounds(collection):
-    n = 0
-    for _ in range(RETRIES):
-        query = {'reducer.status': 'Success'}
-        n = collection.count_documents(query)
-        if n == N_ROUNDS:
-            return n
-        _eprint(f'Succeded rounds {n}. Sleeping for {SLEEP}.')
-        sleep(SLEEP)
-    _eprint(f'Succeded rounds: {n}. Giving up.')
-    return n
-
-if __name__ == '__main__':
-    # Connect to mongo
-    client = pymongo.MongoClient("mongodb://fedn_admin:password@localhost:6534")
-
-    # Wait for successful rounds
-    succeded = _wait_n_rounds(client['fedn-test-network']['control']['round'])
-    assert(succeded == N_ROUNDS) # check that all rounds succeeded
-    _eprint(f'Succeded rounds: {succeded}. Test passed.')
diff --git a/.ci/tests/examples/run.sh b/.ci/tests/examples/run.sh
index 0f887f615..be5650dfe 100755
--- a/.ci/tests/examples/run.sh
+++ b/.ci/tests/examples/run.sh
@@ -17,7 +17,7 @@ docker-compose \
     up -d --build
 
 >&2 echo "Wait for reducer to start"
-sleep 30 # TODO: add API call to check when ready
+".$example/bin/python" ../../.ci/tests/examples/wait_for.py reducer
 
 >&2 echo "Upload compute package"
 curl -k -X POST \
@@ -33,7 +33,7 @@ curl -k -X POST \
 printf '\n'
 
 >&2 echo "Wait for clients to connect"
-".$example/bin/python" ../../.ci/tests/examples/wait_for_clients.py
+".$example/bin/python" ../../.ci/tests/examples/wait_for.py clients
 
 >&2 echo "Start round"
 curl -k -X POST \
@@ -43,7 +43,7 @@ curl -k -X POST \
 printf '\n'
 
 >&2 echo "Checking rounds success"
-".$example/bin/python" ../../.ci/tests/examples/is_success.py
+".$example/bin/python" ../../.ci/tests/examples/wait_for.py rounds
 
 >&2 echo "Test client connection with dowloaded settings"
 # Get config
@@ -57,7 +57,7 @@ docker-compose \
     up -d
 
 >&2 echo "Wait for clients to reconnect"
-".$example/bin/python" ../../.ci/tests/examples/wait_for_clients.py
+".$example/bin/python" ../../.ci/tests/examples/wait_for.py clients
 
 popd
 >&2 echo "Test completed successfully"
\ No newline at end of file
diff --git a/.ci/tests/examples/wait_for.py b/.ci/tests/examples/wait_for.py
new file mode 100644
index 000000000..991ae2e32
--- /dev/null
+++ b/.ci/tests/examples/wait_for.py
@@ -0,0 +1,72 @@
+import json
+import ssl
+import sys
+from time import sleep
+
+import fire
+import pymongo
+import requests
+
+RETRIES = 18
+SLEEP = 10
+
+
+def _eprint(*args, **kwargs):
+    print(*args, file=sys.stderr, **kwargs)
+
+
+def _retry(try_func, **func_args):
+    n = 0
+    for _ in range(RETRIES):
+        is_success = try_func(**func_args)
+        if is_success:
+            _eprint(f'Sucess.')
+            return True
+        _eprint(f'Sleeping for {SLEEP}.')
+        sleep(SLEEP)
+    _eprint(f'Giving up.')
+    return False
+
+
+def _test_rounds(n_rounds):
+    client = pymongo.MongoClient(
+        "mongodb://fedn_admin:password@localhost:6534")
+    collection = client['fedn-test-network']['control']['round']
+    query = {'reducer.status': 'Success'}
+    n = collection.count_documents(query)
+    client.close()
+    _eprint(f'Succeded rounds: {n}.')
+    return n == n_rounds
+
+
+def _test_nodes(n_nodes, node_type, reducer_host='localhost', reducer_port='8090'):
+    try:
+        resp = requests.get(
+            f'https://{reducer_host}:{reducer_port}/netgraph', verify=False)
+    except Exception as e:
+        _eprint(f'Reques exception econuntered: {e}.')
+        return False
+    if resp.status_code == 200:
+        gr = json.loads(resp.content)
+        n = sum(values.get('type') == node_type and values.get(
+            'status') == 'active' for values in gr['nodes'])
+        _eprint(f'Active {node_type}s: {n}.')
+        return n == n_nodes
+    _eprint(f'Reducer returned {resp.status_code}.')
+    return False
+
+
+def rounds(n_rounds=3):
+    _retry(_test_rounds, n_rounds=n_rounds)
+
+
+def clients(n_clients=2):
+    _retry(_test_nodes, n_nodes=n_clients, node_type='client')
+
+
+def reducer():
+    _retry(_test_nodes, n_nodes=1, node_type='reducer')
+
+
+if __name__ == '__main__':
+    fire.Fire()
diff --git a/.ci/tests/examples/wait_for_clients.py b/.ci/tests/examples/wait_for_clients.py
deleted file mode 100644
index 0e5791243..000000000
--- a/.ci/tests/examples/wait_for_clients.py
+++ /dev/null
@@ -1,39 +0,0 @@
-import json
-import sys
-from time import sleep
-
-import pandas
-import requests
-
-N_CLIENTS = 2
-RETRIES = 18
-SLEEP = 10
-REDUCER_HOST = 'localhost'
-
-
-def _eprint(*args, **kwargs):
-    print(*args, file=sys.stderr, **kwargs)
-
-
-def _wait_n_clients():
-    n = 0
-    for _ in range(RETRIES):
-        resp = requests.get(
-            f'https://{REDUCER_HOST}:8090/netgraph', verify=False)
-        if resp.status_code == 200:
-            gr = json.loads(resp.content)
-            n = sum(values.get('type') == 'client' and values.get(
-                'status') == 'active' for values in gr['nodes'])
-            if n == N_CLIENTS:
-                return n
-        _eprint(f'Connected clients: {n}. Sleeping for {SLEEP}.')
-        sleep(SLEEP)
-    _eprint(f'Connected clients: {n}. Giving up.')
-    return n
-
-
-if __name__ == '__main__':
-    # Wait for clients
-    connected = _wait_n_clients()
-    assert(connected == N_CLIENTS)  # check that all clients connected
-    _eprint(f'Connected clients: {connected}. Test passed.')