[test] Make test muting logic more generic

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/backend/BlackBoxCodegenSuppressor.kt

@@ -30,35 +30,15 @@ class BlackBoxCodegenSuppressor(
         val suppressionChecker = testServices.codegenSuppressionChecker
         val moduleStructure = testServices.moduleStructure
         val ignoreDirectives = suppressionChecker.extractIgnoreDirectives(moduleStructure.modules.first()) ?: return failedAssertions
-        for (ignoreDirective in ignoreDirectives) {
+        return suppressionChecker.processAllDirectives(ignoreDirectives) { ignoreDirective, suppressionResult ->
-            val suppressionResult = moduleStructure.modules
+            listOfNotNull(
-                .map { suppressionChecker.failuresInModuleAreIgnored(it, ignoreDirective) }
+                suppressionChecker.processMutedTest(
-                .firstOrNull { it.testMuted }
+                    failed = failedAssertions.isNotEmpty(),
-                ?: continue
+                    ignoreDirective,
-            val additionalMessage = suppressionResult.matchedBackend?.let { "for $it" } ?: ""
+                    suppressionResult,
-            return processAssertions(failedAssertions, ignoreDirective, additionalMessage)
+                )?.wrap()
-        }
+            )
-        return failedAssertions
+        } ?: failedAssertions
-    }
-
-    private fun processAssertions(
-        failedAssertions: List<WrappedException>,
-        directive: ValueDirective<TargetBackend>,
-        additionalMessage: String = ""
-    ): List<WrappedException> {
-        return if (failedAssertions.isNotEmpty()) emptyList()
-        else {
-            val message = buildString {
-                val module = testServices.moduleStructure.modules.first()
-                val targetBackend = testServices.defaultsProvider.defaultTargetBackend ?: module.targetBackend
-                append("Looks like this test can be unmuted. Remove ${targetBackend?.name?.let { "$it from" } ?: "" } ${directive.name} directive for ${module.frontendKind}")
-                if (additionalMessage.isNotEmpty()) {
-                    append(" ")
-                    append(additionalMessage)
-                }
-            }
-            listOf(AssertionError(message).wrap())
-        }
     }
 
     class SuppressionChecker(val testServices: TestServices, val customIgnoreDirective: ValueDirective<TargetBackend>?) : TestService {
@@ -92,6 +72,96 @@ class BlackBoxCodegenSuppressor(
             }
         }
 
+        /**
+         * Finds the first directive from [ignoreDirectives] that mutes this test on the current backend, and
+         * runs [processDirective] for that directive.
+         *
+         * Returns whatever [processDirective] returns, or `null` if this test will not be muted.
+         */
+        @PublishedApi
+        internal inline fun <R> processAllDirectives(
+            ignoreDirectives: List<ValueDirective<TargetBackend>>,
+            processDirective: (ValueDirective<TargetBackend>, SuppressionResult) -> R,
+        ): R? {
+            val modules = testServices.moduleStructure.modules
+            for (ignoreDirective in ignoreDirectives) {
+                val suppressionResult = modules
+                    .map { failuresInModuleAreIgnored(it, ignoreDirective) }
+                    .firstOrNull { it.testMuted }
+                    ?: continue
+                return processDirective(ignoreDirective, suppressionResult)
+            }
+            return null
+        }
+
+        /**
+         * Returns `null` if [failed] is `true`, otherwise returns an [AssertionError] with a message reminding to remove [directive]
+         * from the test to unmute it.
+         */
+        @PublishedApi
+        internal fun processMutedTest(
+            failed: Boolean,
+            directive: ValueDirective<TargetBackend>,
+            suppressionResult: SuppressionResult,
+        ): AssertionError? {
+            if (failed) return null
+
+            val firstModule = testServices.moduleStructure.modules.first()
+            val targetBackend = testServices.defaultsProvider.defaultTargetBackend ?: firstModule.targetBackend
+            val message = buildString {
+                append("Looks like this test can be unmuted. Remove ")
+                targetBackend?.name?.let {
+                    append(it)
+                    append(" from the ")
+                }
+                append(directive.name)
+                append(" directive for ")
+                append(firstModule.frontendKind)
+
+                assert(suppressionResult.testMuted)
+
+                suppressionResult.matchedBackend?.let {
+                    append(" for ")
+                    append(it)
+                }
+            }
+            return AssertionError(message)
+        }
+
+        /**
+         * Runs [block]. If this test has been muted by one of [ignoreDirectives] **and** [block] returns without throwing an exception,
+         * throws an [AssertionError] reminding you to unmute the test.
+         *
+         * If this test has been muted by one of [ignoreDirectives] **and** [block] throws an exception of type [ExpectedError],
+         * catches that exception and returns normally.
+         *
+         * If [block] throws an exception of some other type, rethrows it.
+         *
+         * If this test hasn't been muted **and** [block] throws any exception, rethrows that exception as well.
+         */
+        inline fun <reified ExpectedError : Throwable> checkMuted(
+            ignoreDirectives: List<ValueDirective<TargetBackend>>,
+            block: () -> Unit,
+        ) {
+            val expectedError: ExpectedError? = try {
+                block()
+                null
+            } catch (e: Throwable) {
+                e as? ExpectedError ?: throw e
+            }
+
+            processAllDirectives<Unit>(ignoreDirectives) { ignoreDirective, suppressionResult ->
+                processMutedTest(
+                    failed = expectedError != null,
+                    ignoreDirective,
+                    suppressionResult
+                )?.let { throw it }
+                return
+            }
+
+            expectedError?.let { throw it }
+        }
+
         data class SuppressionResult(val testMuted: Boolean, val matchedBackend: TargetBackend?) {
             companion object {
                 val NO_MUTE = SuppressionResult(false, null)