From 60dc1e411597a83023393945b610fe74d5ff0066 Mon Sep 17 00:00:00 2001
From: Adrian Gruntkowski <adrian.gruntkowski@gmail.com>
Date: Wed, 11 Dec 2024 21:59:57 +0100
Subject: [PATCH] Address flakiness in ingest counter tests

---
 test/plausible/ingestion/counters_test.exs | 14 ++++++++++----
 test/support/test_utils.ex                 |  9 ++++++++-
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/test/plausible/ingestion/counters_test.exs b/test/plausible/ingestion/counters_test.exs
index 5bfaffd163..99e30b0c46 100644
--- a/test/plausible/ingestion/counters_test.exs
+++ b/test/plausible/ingestion/counters_test.exs
@@ -145,14 +145,20 @@ defmodule Plausible.Ingestion.CountersTest do
             r.value == ^value
       )
 
-    query =
+    count =
       if site_id do
-        query |> where([r], r.site_id == ^site_id)
+        query
+        |> where([r], r.site_id == ^site_id)
+        |> await_clickhouse_count(1)
       else
-        query |> where([r], is_nil(r.site_id))
+        # As null site_id entries might be added by parallel tests,
+        # we make the check more permissive in that case.
+        query
+        |> where([r], is_nil(r.site_id))
+        |> await_clickhouse_count(fn count -> count > 0 end)
       end
 
-    assert await_clickhouse_count(query, 1)
+    assert count
   end
 
   defp random_domain() do
diff --git a/test/support/test_utils.ex b/test/support/test_utils.ex
index ea91d8c0c3..ca2aecc49f 100644
--- a/test/support/test_utils.ex
+++ b/test/support/test_utils.ex
@@ -220,11 +220,18 @@ defmodule Plausible.TestUtils do
   end
 
   def await_clickhouse_count(query, expected) do
+    expect_fn =
+      if is_function(expected, 1) do
+        expected
+      else
+        fn count -> count == expected end
+      end
+
     eventually(
       fn ->
         count = Plausible.ClickhouseRepo.aggregate(query, :count)
 
-        {count == expected, count}
+        {expect_fn.(count), count}
       end,
       500,
       10