From 301c8d21dfb4dfa5a12c84430a81e99907288191 Mon Sep 17 00:00:00 2001
From: Ivan Habunek <ivan@habunek.com>
Date: Mon, 1 Jan 2024 11:14:04 +0100
Subject: [PATCH] Add test util function for retrying tests

---
 tests/integration/test_timelines.py | 17 ++++++++---------
 tests/utils.py                      | 23 +++++++++++++++++++++++
 2 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/tests/integration/test_timelines.py b/tests/integration/test_timelines.py
index d4cebbd..1ee3602 100644
--- a/tests/integration/test_timelines.py
+++ b/tests/integration/test_timelines.py
@@ -2,6 +2,7 @@ import pytest
 
 from time import sleep
 from uuid import uuid4
+from tests.utils import run_with_retries
 
 from toot import api, cli
 from toot.entities import from_dict, Status
@@ -40,16 +41,14 @@ def test_timelines(app, user, other_user, friend_user, friend_list, run):
     status2 = _post_status(app, other_user, "#bar")
     status3 = _post_status(app, friend_user, "#foo #bar")
 
-    # Give mastodon time to process things :/
-    # Tests fail if this is removed, required delay depends on server speed
-    sleep(1)
-
     # Home timeline
-    result = run(cli.timelines.timeline)
-    assert result.exit_code == 0
-    assert status1.id in result.stdout
-    assert status2.id not in result.stdout
-    assert status3.id in result.stdout
+    def test_home():
+        result = run(cli.timelines.timeline)
+        assert result.exit_code == 0
+        assert status1.id in result.stdout
+        assert status2.id not in result.stdout
+        assert status3.id in result.stdout
+    run_with_retries(test_home)
 
     # Public timeline
     result = run(cli.timelines.timeline, "--public")
diff --git a/tests/utils.py b/tests/utils.py
index cdae09c..817bdb9 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -2,6 +2,9 @@
 Helpers for testing.
 """
 
+import time
+from typing import Any, Callable
+
 
 class MockResponse:
     def __init__(self, response_data={}, ok=True, is_redirect=False):
@@ -19,3 +22,23 @@ class MockResponse:
 
 def retval(val):
     return lambda *args, **kwargs: val
+
+
+def run_with_retries(fn: Callable[..., Any]):
+    """
+    Run the the given function repeatedly until it finishes without raising an
+    AssertionError. Sleep a bit between attempts. If the function doesn't
+    succeed in the given number of tries raises the AssertionError. Used for
+    tests which should eventually succeed.
+    """
+
+    # Wait upto 6 seconds with incrementally longer sleeps
+    delays = [0.1, 0.2, 0.3, 0.4, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]
+
+    for delay in delays:
+        try:
+            return fn()
+        except AssertionError:
+            time.sleep(delay)
+
+    fn()