diff --git a/src/test/java/com/android/tools/r8/examples/switchmaps/SwitchMapsTestRunner.java b/src/test/java/com/android/tools/r8/examples/switchmaps/SwitchMapsTestRunner.java
index 1fb0770..2936fd8 100644
--- a/src/test/java/com/android/tools/r8/examples/switchmaps/SwitchMapsTestRunner.java
+++ b/src/test/java/com/android/tools/r8/examples/switchmaps/SwitchMapsTestRunner.java
@@ -17,6 +17,34 @@
 @RunWith(Parameterized.class)
 public class SwitchMapsTestRunner extends ExamplesTestBase {
 
+  public static final String EXPECTED =
+      StringUtils.lines(
+          "other",
+          "1, 3 or 4",
+          "2 or 5",
+          "other",
+          "2 or 5",
+          "3 or 5",
+          "1, 3 or 4",
+          "2 or 5",
+          "other",
+          "1, 3 or 4",
+          "2 or 5",
+          "3 or 5",
+          "2 or 5",
+          "other",
+          "6",
+          "7",
+          "7",
+          "rar",
+          "colorful",
+          "blew",
+          "colorful",
+          "soylent",
+          "sooo green",
+          "fifty",
+          "not really");
+
   @Parameterized.Parameters(name = "{0}")
   public static TestParametersCollection data() {
     return getTestParameters().withAllRuntimesAndApiLevels().enableApiLevelsForCf().build();
@@ -42,32 +70,7 @@
 
   @Override
   public String getExpected() {
-    return StringUtils.lines(
-        "other",
-        "1, 3 or 4",
-        "2 or 5",
-        "other",
-        "2 or 5",
-        "3 or 5",
-        "1, 3 or 4",
-        "2 or 5",
-        "other",
-        "1, 3 or 4",
-        "2 or 5",
-        "3 or 5",
-        "2 or 5",
-        "other",
-        "6",
-        "7",
-        "7",
-        "rar",
-        "colorful",
-        "blew",
-        "colorful",
-        "soylent",
-        "sooo green",
-        "fifty",
-        "not really");
+    return EXPECTED;
   }
 
   @Test
diff --git a/src/test/java/com/android/tools/r8/rewrite/switchmaps/RewriteSwitchMapsTest.java b/src/test/java/com/android/tools/r8/rewrite/switchmaps/RewriteSwitchMapsTest.java
index 98d8604..ab4a619 100644
--- a/src/test/java/com/android/tools/r8/rewrite/switchmaps/RewriteSwitchMapsTest.java
+++ b/src/test/java/com/android/tools/r8/rewrite/switchmaps/RewriteSwitchMapsTest.java
@@ -4,16 +4,17 @@
 package com.android.tools.r8.rewrite.switchmaps;
 
 import static com.android.tools.r8.utils.codeinspector.Matchers.isAbsent;
+import static com.android.tools.r8.utils.codeinspector.Matchers.isPresent;
 import static org.hamcrest.MatcherAssert.assertThat;
 
 import com.android.tools.r8.TestBase;
 import com.android.tools.r8.TestParameters;
 import com.android.tools.r8.TestParametersCollection;
-import com.android.tools.r8.ToolHelper;
+import com.android.tools.r8.examples.switchmaps.Colors;
+import com.android.tools.r8.examples.switchmaps.Days;
+import com.android.tools.r8.examples.switchmaps.SwitchMapsTestRunner;
+import com.android.tools.r8.examples.switchmaps.Switches;
 import com.android.tools.r8.utils.AndroidApiLevel;
-import com.google.common.collect.ImmutableList;
-import java.nio.file.Paths;
-import java.util.List;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
@@ -31,20 +32,29 @@
     return getTestParameters().withDefaultRuntimes().withApiLevel(AndroidApiLevel.B).build();
   }
 
-  private static final String JAR_FILE = "switchmaps.jar";
-  private static final String SWITCHMAP_CLASS_NAME = "switchmaps.Switches$1";
-  private static final List<String> PG_CONFIG = ImmutableList.of(
-      "-keep class switchmaps.Switches { public static void main(...); }",
-      "-dontobfuscate",
-      "-keepattributes *");
+  private static final String SWITCHMAP_CLASS_NAME = typeName(Switches.class) + "$1";
+
+  @Test
+  public void testReference() throws Exception {
+    testForRuntime(parameters)
+        .addProgramClassesAndInnerClasses(Switches.class)
+        .addProgramClasses(Colors.class, Days.class)
+        .run(parameters.getRuntime(), Switches.class)
+        .assertSuccessWithOutput(SwitchMapsTestRunner.EXPECTED)
+        .inspect(inspector -> assertThat(inspector.clazz(SWITCHMAP_CLASS_NAME), isPresent()));
+  }
 
   @Test
   public void checkSwitchMapsRemoved() throws Exception {
     testForR8(parameters.getBackend())
-        .addProgramFiles(Paths.get(ToolHelper.EXAMPLES_BUILD_DIR).resolve(JAR_FILE))
-        .addKeepRules(PG_CONFIG)
+        .addProgramClassesAndInnerClasses(Switches.class)
+        .addProgramClasses(Colors.class, Days.class)
+        .addKeepMainRule(Switches.class)
+        .addDontObfuscate()
+        .addKeepAllAttributes()
         .setMinApi(parameters)
-        .compile()
+        .run(parameters.getRuntime(), Switches.class)
+        .assertSuccessWithOutput(SwitchMapsTestRunner.EXPECTED)
         .inspect(inspector -> assertThat(inspector.clazz(SWITCHMAP_CLASS_NAME), isAbsent()));
   }
 }