1 files changed, 6 insertions, 18 deletions

diff --git a/spark-common/src/test/java/me/lucko/spark/common/sampler/SamplerTest.java b/spark-common/src/test/java/me/lucko/spark/common/sampler/SamplerTest.java
index f6d9f87..53fc7e6 100644
--- a/spark-common/src/test/java/me/lucko/spark/common/sampler/SamplerTest.java
+++ b/spark-common/src/test/java/me/lucko/spark/common/sampler/SamplerTest.java
@@ -20,8 +20,6 @@
 
 package me.lucko.spark.common.sampler;
 
-import me.lucko.spark.common.sampler.async.AsyncSampler;
-import me.lucko.spark.common.sampler.java.JavaSampler;
 import me.lucko.spark.common.sampler.java.MergeStrategy;
 import me.lucko.spark.common.sampler.source.ClassSourceLookup;
 import me.lucko.spark.proto.SparkSamplerProtos;
@@ -48,8 +46,8 @@ public class SamplerTest {
 
     @ParameterizedTest
     @EnumSource
-    public void testSampler(ProfilerType profilerType, @TempDir Path directory) {
-        if (profilerType == ProfilerType.ASYNC) {
+    public void testSampler(SamplerType samplerType, @TempDir Path directory) {
+        if (samplerType == SamplerType.ASYNC) {
             String os = System.getProperty("os.name").toLowerCase(Locale.ROOT).replace(" ", "");
             assumeTrue(os.equals("linux") || os.equals("macosx"), "async profiler is only supported on Linux and macOS");
         }
@@ -62,11 +60,12 @@ public class SamplerTest {
                     .threadDumper(new ThreadDumper.Specific(thread))
                     .threadGrouper(ThreadGrouper.BY_POOL)
                     .samplingInterval(10)
-                    .forceJavaSampler(profilerType == ProfilerType.JAVA)
+                    .forceJavaSampler(samplerType == SamplerType.JAVA)
                     .completeAfter(2, TimeUnit.SECONDS)
                     .start(plugin.platform());
 
-            assertInstanceOf(profilerType.expectedClass, sampler);
+            assertInstanceOf(samplerType.implClass(), sampler);
+            assertEquals(samplerType, sampler.getType());
 
             assertNotEquals(-1, sampler.getAutoEndTime());
             sampler.getFuture().join();
@@ -75,7 +74,7 @@ public class SamplerTest {
                     .creator(TestCommandSender.INSTANCE.toData())
                     .classSourceLookup(() -> ClassSourceLookup.create(plugin.platform()));
 
-            if (profilerType == ProfilerType.JAVA) {
+            if (samplerType == SamplerType.JAVA) {
                 exportProps.mergeStrategy(MergeStrategy.SAME_METHOD);
             }
 
@@ -93,15 +92,4 @@ public class SamplerTest {
         }
     }
 
-    public enum ProfilerType {
-        JAVA(JavaSampler.class),
-        ASYNC(AsyncSampler.class);
-
-        private final Class<? extends Sampler> expectedClass;
-
-        ProfilerType(Class<? extends Sampler> expectedClass) {
-            this.expectedClass = expectedClass;
-        }
-    }
-
 }