From 630875dae23d9c2adc2920fc7f990613cbe2b92f Mon Sep 17 00:00:00 2001
From: Greyson Parrelli <greyson@signal.org>
Date: Thu, 15 Aug 2024 09:29:43 -0400
Subject: [PATCH] Reduce noise of flaky test.

---
 .../backup/v2/FlakyTestAnnotationTest.kt      | 32 ++++++++++++
 .../securesms/database/SQLiteDatabaseTest.kt  |  7 +++
 .../securesms/testing/SignalFlakyTest.kt      |  9 ++++
 .../securesms/testing/SignalFlakyTestRule.kt  | 49 +++++++++++++++++++
 4 files changed, 97 insertions(+)
 create mode 100644 app/src/androidTest/java/org/thoughtcrime/securesms/backup/v2/FlakyTestAnnotationTest.kt
 create mode 100644 app/src/androidTest/java/org/thoughtcrime/securesms/testing/SignalFlakyTest.kt
 create mode 100644 app/src/androidTest/java/org/thoughtcrime/securesms/testing/SignalFlakyTestRule.kt

diff --git a/app/src/androidTest/java/org/thoughtcrime/securesms/backup/v2/FlakyTestAnnotationTest.kt b/app/src/androidTest/java/org/thoughtcrime/securesms/backup/v2/FlakyTestAnnotationTest.kt
new file mode 100644
index 0000000000..3768da67eb
--- /dev/null
+++ b/app/src/androidTest/java/org/thoughtcrime/securesms/backup/v2/FlakyTestAnnotationTest.kt
@@ -0,0 +1,32 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.backup.v2
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import org.junit.Assert.assertEquals
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.thoughtcrime.securesms.testing.SignalFlakyTest
+import org.thoughtcrime.securesms.testing.SignalFlakyTestRule
+
+@RunWith(AndroidJUnit4::class)
+class FlakyTestAnnotationTest {
+
+  @get:Rule
+  val flakyTestRule = SignalFlakyTestRule()
+
+  companion object {
+    private var count = 0
+  }
+
+  @SignalFlakyTest
+  @Test
+  fun purposelyFlaky() {
+    count++
+    assertEquals(3, count)
+  }
+}
diff --git a/app/src/androidTest/java/org/thoughtcrime/securesms/database/SQLiteDatabaseTest.kt b/app/src/androidTest/java/org/thoughtcrime/securesms/database/SQLiteDatabaseTest.kt
index ba6d6aa712..b532125504 100644
--- a/app/src/androidTest/java/org/thoughtcrime/securesms/database/SQLiteDatabaseTest.kt
+++ b/app/src/androidTest/java/org/thoughtcrime/securesms/database/SQLiteDatabaseTest.kt
@@ -4,9 +4,12 @@ import androidx.test.ext.junit.runners.AndroidJUnit4
 import junit.framework.Assert.assertFalse
 import junit.framework.Assert.assertTrue
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.signal.core.util.concurrent.SignalExecutors
+import org.thoughtcrime.securesms.testing.SignalFlakyTest
+import org.thoughtcrime.securesms.testing.SignalFlakyTestRule
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.atomic.AtomicBoolean
 
@@ -18,6 +21,9 @@ class SQLiteDatabaseTest {
 
   private lateinit var db: SQLiteDatabase
 
+  @get:Rule
+  val flakyTestRule = SignalFlakyTestRule()
+
   @Before
   fun setup() {
     db = SignalDatabase.instance!!.signalWritableDatabase
@@ -181,6 +187,7 @@ class SQLiteDatabaseTest {
     assertTrue(hasRun2.get())
   }
 
+  @SignalFlakyTest
   @Test
   fun runPostSuccessfulTransaction_runsAfterMainTransactionInNestedTransaction() {
     val hasRun1 = AtomicBoolean(false)
diff --git a/app/src/androidTest/java/org/thoughtcrime/securesms/testing/SignalFlakyTest.kt b/app/src/androidTest/java/org/thoughtcrime/securesms/testing/SignalFlakyTest.kt
new file mode 100644
index 0000000000..254752d3f5
--- /dev/null
+++ b/app/src/androidTest/java/org/thoughtcrime/securesms/testing/SignalFlakyTest.kt
@@ -0,0 +1,9 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.testing
+
+@Retention(AnnotationRetention.RUNTIME)
+annotation class SignalFlakyTest(val allowedAttempts: Int = 3)
diff --git a/app/src/androidTest/java/org/thoughtcrime/securesms/testing/SignalFlakyTestRule.kt b/app/src/androidTest/java/org/thoughtcrime/securesms/testing/SignalFlakyTestRule.kt
new file mode 100644
index 0000000000..b12e18e174
--- /dev/null
+++ b/app/src/androidTest/java/org/thoughtcrime/securesms/testing/SignalFlakyTestRule.kt
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.testing
+
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+import org.signal.core.util.logging.Log
+
+/**
+ * A JUnit rule that retries tests annotated with [SignalFlakyTest] before considering them to be a failure.
+ * As the name implies, this is useful for known-flaky tests.
+ */
+class SignalFlakyTestRule : TestRule {
+  override fun apply(base: Statement, description: Description): Statement {
+    val flakyAnnotation = description.getAnnotation(SignalFlakyTest::class.java)
+
+    return if (flakyAnnotation != null) {
+      FlakyStatement(
+        base = base,
+        description = description,
+        allowedAttempts = flakyAnnotation.allowedAttempts
+      )
+    } else {
+      base
+    }
+  }
+
+  private class FlakyStatement(private val base: Statement, private val description: Description, private val allowedAttempts: Int) : Statement() {
+    override fun evaluate() {
+      var attemptsRemaining = allowedAttempts
+      while (attemptsRemaining > 0) {
+        try {
+          base.evaluate()
+          return
+        } catch (t: Throwable) {
+          attemptsRemaining--
+          if (attemptsRemaining <= 0) {
+            throw t
+          }
+          Log.w(description.testClass.simpleName, "[${description.methodName}] Flaky test failed! $attemptsRemaining attempt(s) remaining.", t)
+        }
+      }
+    }
+  }
+}