From fbb8c1be9bfe78ae3dbbe04b2d5b1d32fbee89a2 Mon Sep 17 00:00:00 2001
From: Shumaf Lovpache <shumaf.lovpache@jetbrains.com>
Date: Tue, 14 Jan 2025 18:52:20 +0000
Subject: [PATCH] [platform test framework] Change the registry value type to
 String

Merge-request: IJ-MR-152137
Merged-by: Shumaf Lovpache <soarex16@gmail.com>

GitOrigin-RevId: 9826a6f8794a042e7f8e9067629ae381b14cbae3
---
 .../testFramework/junit5/src/RegistryKey.kt   | 16 +++++++
 .../junit5/src/RegistryKeyExtension.kt        | 34 --------------
 .../junit5/src/impl/RegistryKeyExtension.kt   | 47 +++++++++++++++++++
 3 files changed, 63 insertions(+), 34 deletions(-)
 create mode 100644 platform/testFramework/junit5/src/RegistryKey.kt
 delete mode 100644 platform/testFramework/junit5/src/RegistryKeyExtension.kt
 create mode 100644 platform/testFramework/junit5/src/impl/RegistryKeyExtension.kt

diff --git a/platform/testFramework/junit5/src/RegistryKey.kt b/platform/testFramework/junit5/src/RegistryKey.kt
new file mode 100644
index 000000000000..8ecd591b2613
--- /dev/null
+++ b/platform/testFramework/junit5/src/RegistryKey.kt
@@ -0,0 +1,16 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.testFramework.junit5
+
+import com.intellij.testFramework.junit5.impl.RegistryKeyExtension
+import org.jetbrains.annotations.TestOnly
+import org.junit.jupiter.api.extension.ExtendWith
+
+/**
+ * Sets the [value] for the Registry [key], runs, and reverts to previous value.
+ * @see com.intellij.openapi.util.registry.RegistryValue.setValue
+ */
+@TestOnly
+@Repeatable
+@Target(AnnotationTarget.FUNCTION, AnnotationTarget.CLASS)
+@ExtendWith(RegistryKeyExtension::class)
+annotation class RegistryKey(val key: String, val value: String)
\ No newline at end of file
diff --git a/platform/testFramework/junit5/src/RegistryKeyExtension.kt b/platform/testFramework/junit5/src/RegistryKeyExtension.kt
deleted file mode 100644
index 23f16295d5d9..000000000000
--- a/platform/testFramework/junit5/src/RegistryKeyExtension.kt
+++ /dev/null
@@ -1,34 +0,0 @@
-// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
-package com.intellij.testFramework.junit5
-
-import com.intellij.openapi.util.registry.Registry
-import org.junit.jupiter.api.extension.AfterEachCallback
-import org.junit.jupiter.api.extension.BeforeEachCallback
-import org.junit.jupiter.api.extension.ExtensionContext
-
-/**
- * Sets the [value] for the Registry [key], runs, and reverts to previous value.
- *
- * Boolean keys are supported currently.
- * TODO change the [value] type to [String] when needed.
- *
- * @see com.intellij.openapi.util.registry.RegistryValue.setValue
- */
-class RegistryKeyExtension(
-  key: String,
-  private val value: Boolean,
-) : BeforeEachCallback,
-    AfterEachCallback {
-
-  private val registryValue = Registry.get(key)
-  private var previous: Boolean = false
-
-  override fun beforeEach(context: ExtensionContext?) {
-    previous = registryValue.asBoolean()
-    registryValue.setValue(value)
-  }
-
-  override fun afterEach(context: ExtensionContext?) {
-    registryValue.setValue(previous)
-  }
-}
diff --git a/platform/testFramework/junit5/src/impl/RegistryKeyExtension.kt b/platform/testFramework/junit5/src/impl/RegistryKeyExtension.kt
new file mode 100644
index 000000000000..686b21d61f15
--- /dev/null
+++ b/platform/testFramework/junit5/src/impl/RegistryKeyExtension.kt
@@ -0,0 +1,47 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.testFramework.junit5.impl
+
+import com.intellij.openapi.util.registry.Registry
+import com.intellij.testFramework.junit5.RegistryKey
+import org.jetbrains.annotations.TestOnly
+import org.junit.jupiter.api.extension.ExtensionContext
+import org.junit.jupiter.api.extension.InvocationInterceptor
+import org.junit.platform.commons.support.AnnotationSupport
+
+@TestOnly
+internal class RegistryKeyExtension : AbstractInvocationInterceptor() {
+
+  override fun <T> intercept(invocation: InvocationInterceptor.Invocation<T>, context: ExtensionContext): T {
+    val annotations = AnnotationSupport.findRepeatableAnnotations(context.testClass, RegistryKey::class.java) +
+                      AnnotationSupport.findRepeatableAnnotations(context.element, RegistryKey::class.java)
+    if (annotations.isEmpty()) {
+      return invocation.proceed()
+    }
+    val valuesBefore = annotations.map { annotation ->
+      Pair(annotation, annotation.setValue())
+    }
+    try {
+      return invocation.proceed()
+    }
+    finally {
+      for ((annotation, previousValue) in valuesBefore.asReversed()) {
+        annotation.resetValue(previousValue)
+      }
+    }
+  }
+
+  private fun RegistryKey.setValue(): String? = Registry.Companion.get(key).run {
+    val previousValue = Registry.Companion.stringValue(key)
+    setValue(value)
+    previousValue
+  }
+
+  private fun RegistryKey.resetValue(previousValue: String?): Unit = Registry.Companion.get(key).run {
+    if (previousValue == null) {
+      resetToDefault()
+    }
+    else {
+      setValue(previousValue)
+    }
+  }
+}
\ No newline at end of file