Add unit tests for SnapshotsPreviewRuntimeRetentionTransform

gradle-plugin/plugin/build.gradle.kts

@@ -27,6 +27,9 @@ java {
 				srcDir(perfProjectTemplateResDir)
 			}
 		}
+    test {
+      java.srcDir("src/test/kotlin")
+    }
 	}
 }
 
@@ -99,10 +102,6 @@ tasks.check {
 	dependsOn(functionalTestTask)
 }
 
-tasks.withType<Test> {
-	useJUnitPlatform()
-}
-
 val packagePerformanceProjectTemplateTask =
 	tasks.register<Zip>("packagePerformanceProjectTemplate") {
 		archiveFileName.set("performance-project-template.zip")
@@ -134,6 +133,7 @@ dependencies {
   // from one dependency conflicting with that of dexlib2, so we'll use the same version here.
   implementation(libs.guava)
 
+  testImplementation(libs.junit)
 	testImplementation(libs.junit.jupiter.api)
 	testImplementation(libs.google.truth)
 	testRuntimeOnly(libs.junit.jupiter.engine)

gradle-plugin/plugin/src/test/kotlin/com/emergetools/android/gradle/instrumentation/snapshots/SnapshotsPreviewRuntimeRetentionTransformTest.kt

@@ -0,0 +1,131 @@
+package com.emergetools.android.gradle.instrumentation.snapshots
+
+import com.emergetools.android.gradle.instrumentation.testutils.isClassAnnotationRuntimeVisible
+import com.emergetools.android.gradle.instrumentation.testutils.isMethodAnnotationRuntimeVisible
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertTrue
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+import org.objectweb.asm.ClassReader
+import org.objectweb.asm.ClassWriter
+import org.objectweb.asm.Opcodes
+import org.slf4j.LoggerFactory
+import java.io.File
+import java.io.InputStream
+
+@RunWith(JUnit4::class)
+class SnapshotsPreviewRuntimeRetentionTransformTest {
+
+  @Test
+  fun `test transform applied correctly on method with preview annotation`() {
+    val originalClassBytes = loadClassFile("snapshot-test-classes/TextRowWithIconKt.class")
+
+    // Ensure annotation is not yet runtime visible
+    assertFalse(
+      isMethodAnnotationRuntimeVisible(
+        originalClassBytes,
+        "Landroidx/compose/ui/tooling/preview/Preview;",
+        "TextRowWithIconPreviewFromMainIgnored",
+      )
+    )
+
+    val transformedClassBytes = applyTransform(originalClassBytes)
+
+    assertTrue(
+      isMethodAnnotationRuntimeVisible(
+        transformedClassBytes,
+        "Landroidx/compose/ui/tooling/preview/Preview;",
+        "TextRowWithIconPreviewFromMainIgnored",
+      )
+    )
+  }
+
+  @Test
+  fun `test transform applied correctly on method with preview container annotation`() {
+    val originalClassBytes = loadClassFile("snapshot-test-classes/TextRowWithIconKt.class")
+
+    // Ensure annotation is not yet runtime visible
+    assertFalse(
+      isMethodAnnotationRuntimeVisible(
+        originalClassBytes,
+        "Landroidx/compose/ui/tooling/preview/Preview\$Container;",
+        "TextRowWithIconPreviewFromMain",
+      )
+    )
+
+    val transformedClassBytes = applyTransform(originalClassBytes)
+
+    assertTrue(
+      isMethodAnnotationRuntimeVisible(
+        transformedClassBytes,
+        "Landroidx/compose/ui/tooling/preview/Preview\$Container;",
+        "TextRowWithIconPreviewFromMain",
+      )
+    )
+  }
+
+  @Test
+  fun `test transform applied correctly on multipreview class with preview annotation`() {
+    val originalClassBytes =
+      loadClassFile("snapshot-test-classes/TestSinglePreviewMultiPreview.class")
+
+    // Ensure annotation is not yet runtime visible
+    assertFalse(
+      isClassAnnotationRuntimeVisible(
+        originalClassBytes,
+        "Landroidx/compose/ui/tooling/preview/Preview;",
+      )
+    )
+
+    val transformedClassBytes = applyTransform(originalClassBytes)
+
+    assertTrue(
+      isClassAnnotationRuntimeVisible(
+        transformedClassBytes,
+        "Landroidx/compose/ui/tooling/preview/Preview;",
+      )
+    )
+  }
+
+  @Test
+  fun `test transform applied correctly on multipreview class with preview container annotation`() {
+    val originalClassBytes = loadClassFile("snapshot-test-classes/TestMultiPreview.class")
+
+    // Ensure annotation is not yet runtime visible
+    assertFalse(
+      isClassAnnotationRuntimeVisible(
+        originalClassBytes,
+        "Landroidx/compose/ui/tooling/preview/Preview\$Container;",
+      )
+    )
+
+    val transformedClassBytes = applyTransform(originalClassBytes)
+
+    assertTrue(
+      isClassAnnotationRuntimeVisible(
+        transformedClassBytes,
+        "Landroidx/compose/ui/tooling/preview/Preview\$Container;",
+      )
+    )
+  }
+
+  private fun loadClassFile(fileName: String): ByteArray {
+    val classFile = File(javaClass.getResource("/$fileName").file)
+    val inputStream: InputStream = classFile.inputStream()
+    return inputStream.readBytes()
+  }
+
+  private fun applyTransform(originalClassBytes: ByteArray): ByteArray {
+    val classReader = ClassReader(originalClassBytes)
+    val classWriter = ClassWriter(classReader, 0)
+
+    val classVisitor = SnapshotsPreviewRuntimeRetentionTransform(
+      Opcodes.ASM9,
+      classWriter,
+      LoggerFactory.getLogger(SnapshotsPreviewRuntimeRetentionTransformTest::class.java)
+    )
+    classReader.accept(classVisitor, ClassReader.SKIP_FRAMES)
+    return classWriter.toByteArray()
+  }
+}

gradle-plugin/plugin/src/test/kotlin/com/emergetools/android/gradle/instrumentation/testutils/TestAnnotationCollector.kt

@@ -0,0 +1,61 @@
+package com.emergetools.android.gradle.instrumentation.testutils
+
+import org.objectweb.asm.AnnotationVisitor
+import org.objectweb.asm.ClassReader
+import org.objectweb.asm.ClassVisitor
+import org.objectweb.asm.MethodVisitor
+import org.objectweb.asm.Opcodes
+
+fun isClassAnnotationRuntimeVisible(
+  classBytes: ByteArray,
+  annotationDescriptor: String,
+): Boolean {
+  val classReader = ClassReader(classBytes)
+  val annotationCollector = TestAnnotationCollector()
+  classReader.accept(annotationCollector, ClassReader.SKIP_FRAMES)
+  return annotationCollector.classAnnotations[annotationDescriptor] ?: false
+}
+
+fun isMethodAnnotationRuntimeVisible(
+  classBytes: ByteArray,
+  annotationDescriptor: String,
+  methodName: String,
+): Boolean {
+  val classReader = ClassReader(classBytes)
+  val annotationCollector = TestAnnotationCollector()
+  classReader.accept(annotationCollector, ClassReader.SKIP_FRAMES)
+  return annotationCollector.methodAnnotations[methodName]?.get(annotationDescriptor) ?: false
+}
+
+private class TestAnnotationCollector : ClassVisitor(Opcodes.ASM9) {
+  val classAnnotations = mutableMapOf<String, Boolean>()
+  val methodAnnotations = mutableMapOf<String, MutableMap<String, Boolean>>()
+
+  override fun visitAnnotation(
+    descriptor: String?,
+    visible: Boolean,
+  ): AnnotationVisitor? {
+    classAnnotations[descriptor!!] = visible
+    return super.visitAnnotation(descriptor, visible)
+  }
+
+  override fun visitMethod(
+    access: Int,
+    name: String,
+    descriptor: String?,
+    signature: String?,
+    exceptions: Array<out String>?,
+  ): MethodVisitor {
+    methodAnnotations[name] = mutableMapOf()
+
+    return object : MethodVisitor(Opcodes.ASM9) {
+      override fun visitAnnotation(
+        descriptor: String,
+        visible: Boolean,
+      ): AnnotationVisitor? {
+        methodAnnotations[name]?.put(descriptor, visible)
+        return super.visitAnnotation(descriptor, visible)
+      }
+    }
+  }
+}