Cleanup of redefinition testing

Move redefine logic into a single common function and perform some
other cleanup.

Test: mma -j40 test-art-host
Change-Id: I8618bda4f392b683ce198374066e356b87578e7b
diff --git a/runtime/openjdkjvmti/OpenjdkJvmTi.cc b/runtime/openjdkjvmti/OpenjdkJvmTi.cc
index d1c2293..1ad3f08 100644
--- a/runtime/openjdkjvmti/OpenjdkJvmTi.cc
+++ b/runtime/openjdkjvmti/OpenjdkJvmTi.cc
@@ -1142,6 +1142,34 @@
     return RetransformClassesWithHook(reinterpret_cast<ArtJvmTiEnv*>(env), classes, hook);
   }
 
+  static jvmtiError RedefineClassDirect(ArtJvmTiEnv* env,
+                                        jclass klass,
+                                        jint dex_size,
+                                        unsigned char* dex_file) {
+    if (!IsValidEnv(env)) {
+      return ERR(INVALID_ENVIRONMENT);
+    }
+    jvmtiError ret = OK;
+    std::string location;
+    if ((ret = GetClassLocation(env, klass, &location)) != OK) {
+      // TODO Do something more here? Maybe give log statements?
+      return ret;
+    }
+    std::string error;
+    ret = Redefiner::RedefineClass(env,
+                                    art::Runtime::Current(),
+                                    art::Thread::Current(),
+                                    klass,
+                                    location,
+                                    dex_size,
+                                    reinterpret_cast<uint8_t*>(dex_file),
+                                    &error);
+    if (ret != OK) {
+      LOG(ERROR) << "FAILURE TO REDEFINE " << error;
+    }
+    return ret;
+  }
+
   // TODO This will be called by the event handler for the art::ti Event Load Event
   static jvmtiError RetransformClassesWithHook(ArtJvmTiEnv* env,
                                                const std::vector<jclass>& classes,
@@ -1252,7 +1280,10 @@
   reinterpret_cast<void*>(JvmtiFunctions::RetransformClassWithHook),
   // nullptr,  // reserved1
   JvmtiFunctions::SetEventNotificationMode,
-  nullptr,  // reserved3
+  // SPECIAL FUNCTION: RedefineClassDirect Is normally reserved3
+  // TODO Remove once we have events working.
+  reinterpret_cast<void*>(JvmtiFunctions::RedefineClassDirect),
+  // nullptr,  // reserved3
   JvmtiFunctions::GetAllThreads,
   JvmtiFunctions::SuspendThread,
   JvmtiFunctions::ResumeThread,
diff --git a/runtime/openjdkjvmti/transform.cc b/runtime/openjdkjvmti/transform.cc
index f7b8b92..f545125 100644
--- a/runtime/openjdkjvmti/transform.cc
+++ b/runtime/openjdkjvmti/transform.cc
@@ -58,6 +58,21 @@
 
 namespace openjdkjvmti {
 
+jvmtiError GetClassLocation(ArtJvmTiEnv* env, jclass klass, /*out*/std::string* location) {
+  JNIEnv* jni_env = nullptr;
+  jint ret = env->art_vm->GetEnv(reinterpret_cast<void**>(&jni_env), JNI_VERSION_1_1);
+  if (ret != JNI_OK) {
+    // TODO Different error might be better?
+    return ERR(INTERNAL);
+  }
+  art::ScopedObjectAccess soa(jni_env);
+  art::StackHandleScope<1> hs(art::Thread::Current());
+  art::Handle<art::mirror::Class> hs_klass(hs.NewHandle(soa.Decode<art::mirror::Class>(klass)));
+  const art::DexFile& dex = hs_klass->GetDexFile();
+  *location = dex.GetLocation();
+  return OK;
+}
+
 // TODO Move this function somewhere more appropriate.
 // Gets the data surrounding the given class.
 jvmtiError GetTransformationData(ArtJvmTiEnv* env,
diff --git a/runtime/openjdkjvmti/transform.h b/runtime/openjdkjvmti/transform.h
index 35b990b..0ad5099 100644
--- a/runtime/openjdkjvmti/transform.h
+++ b/runtime/openjdkjvmti/transform.h
@@ -41,6 +41,8 @@
 
 namespace openjdkjvmti {
 
+jvmtiError GetClassLocation(ArtJvmTiEnv* env, jclass klass, /*out*/std::string* location);
+
 // Gets the data surrounding the given class.
 jvmtiError GetTransformationData(ArtJvmTiEnv* env,
                                  jclass klass,
diff --git a/test/902-hello-transformation/expected.txt b/test/902-hello-transformation/expected.txt
index a826f93..4774b81 100644
--- a/test/902-hello-transformation/expected.txt
+++ b/test/902-hello-transformation/expected.txt
@@ -1,3 +1,2 @@
 hello
-modifying class 'Transform'
 Goodbye
diff --git a/test/902-hello-transformation/run b/test/902-hello-transformation/run
index 3755d1d..94a8b2d 100755
--- a/test/902-hello-transformation/run
+++ b/test/902-hello-transformation/run
@@ -27,12 +27,11 @@
   arg="jvm"
 else
   arg="art"
-fi
-
-if [[ "$@" != *"--debuggable"* ]]; then
-  other_args=" -Xcompiler-option --debuggable "
-else
-  other_args=""
+  if [[ "$@" != *"--debuggable"* ]]; then
+    other_args=" -Xcompiler-option --debuggable "
+  else
+    other_args=""
+  fi
 fi
 
 ./default-run "$@" --experimental agents \
diff --git a/test/902-hello-transformation/src/Main.java b/test/902-hello-transformation/src/Main.java
index 204b6e7..ec47119 100644
--- a/test/902-hello-transformation/src/Main.java
+++ b/test/902-hello-transformation/src/Main.java
@@ -14,7 +14,40 @@
  * limitations under the License.
  */
 
+import java.util.Base64;
 public class Main {
+
+  /**
+   * base64 encoded class/dex file for
+   * class Transform {
+   *   public void sayHi() {
+   *    System.out.println("Goodbye");
+   *   }
+   * }
+   */
+  private static final byte[] CLASS_BYTES = Base64.getDecoder().decode(
+    "yv66vgAAADQAHAoABgAOCQAPABAIABEKABIAEwcAFAcAFQEABjxpbml0PgEAAygpVgEABENvZGUB" +
+    "AA9MaW5lTnVtYmVyVGFibGUBAAVzYXlIaQEAClNvdXJjZUZpbGUBAA5UcmFuc2Zvcm0uamF2YQwA" +
+    "BwAIBwAWDAAXABgBAAdHb29kYnllBwAZDAAaABsBAAlUcmFuc2Zvcm0BABBqYXZhL2xhbmcvT2Jq" +
+    "ZWN0AQAQamF2YS9sYW5nL1N5c3RlbQEAA291dAEAFUxqYXZhL2lvL1ByaW50U3RyZWFtOwEAE2ph" +
+    "dmEvaW8vUHJpbnRTdHJlYW0BAAdwcmludGxuAQAVKExqYXZhL2xhbmcvU3RyaW5nOylWACAABQAG" +
+    "AAAAAAACAAAABwAIAAEACQAAAB0AAQABAAAABSq3AAGxAAAAAQAKAAAABgABAAAAEQABAAsACAAB" +
+    "AAkAAAAlAAIAAQAAAAmyAAISA7YABLEAAAABAAoAAAAKAAIAAAATAAgAFAABAAwAAAACAA0=");
+  private static final byte[] DEX_BYTES = Base64.getDecoder().decode(
+    "ZGV4CjAzNQCLXSBQ5FiS3f16krSYZFF8xYZtFVp0GRXMAgAAcAAAAHhWNBIAAAAAAAAAACwCAAAO" +
+    "AAAAcAAAAAYAAACoAAAAAgAAAMAAAAABAAAA2AAAAAQAAADgAAAAAQAAAAABAACsAQAAIAEAAGIB" +
+    "AABqAQAAcwEAAIABAACXAQAAqwEAAL8BAADTAQAA4wEAAOYBAADqAQAA/gEAAAMCAAAMAgAAAgAA" +
+    "AAMAAAAEAAAABQAAAAYAAAAIAAAACAAAAAUAAAAAAAAACQAAAAUAAABcAQAABAABAAsAAAAAAAAA" +
+    "AAAAAAAAAAANAAAAAQABAAwAAAACAAAAAAAAAAAAAAAAAAAAAgAAAAAAAAAHAAAAAAAAAB4CAAAA" +
+    "AAAAAQABAAEAAAATAgAABAAAAHAQAwAAAA4AAwABAAIAAAAYAgAACQAAAGIAAAAbAQEAAABuIAIA" +
+    "EAAOAAAAAQAAAAMABjxpbml0PgAHR29vZGJ5ZQALTFRyYW5zZm9ybTsAFUxqYXZhL2lvL1ByaW50" +
+    "U3RyZWFtOwASTGphdmEvbGFuZy9PYmplY3Q7ABJMamF2YS9sYW5nL1N0cmluZzsAEkxqYXZhL2xh" +
+    "bmcvU3lzdGVtOwAOVHJhbnNmb3JtLmphdmEAAVYAAlZMABJlbWl0dGVyOiBqYWNrLTMuMzYAA291" +
+    "dAAHcHJpbnRsbgAFc2F5SGkAEQAHDgATAAcOhQAAAAEBAICABKACAQG4Ag0AAAAAAAAAAQAAAAAA" +
+    "AAABAAAADgAAAHAAAAACAAAABgAAAKgAAAADAAAAAgAAAMAAAAAEAAAAAQAAANgAAAAFAAAABAAA" +
+    "AOAAAAAGAAAAAQAAAAABAAABIAAAAgAAACABAAABEAAAAQAAAFwBAAACIAAADgAAAGIBAAADIAAA" +
+    "AgAAABMCAAAAIAAAAQAAAB4CAAAAEAAAAQAAACwCAAA=");
+
   public static void main(String[] args) {
     System.loadLibrary(args[1]);
     doTest(new Transform());
@@ -22,10 +55,12 @@
 
   public static void doTest(Transform t) {
     t.sayHi();
-    doClassTransformation(Transform.class);
+    doCommonClassRedefinition(Transform.class, CLASS_BYTES, DEX_BYTES);
     t.sayHi();
   }
 
   // Transforms the class
-  private static native void doClassTransformation(Class target);
+  private static native void doCommonClassRedefinition(Class<?> target,
+                                                       byte[] class_file,
+                                                       byte[] dex_file);
 }
diff --git a/test/902-hello-transformation/transform.cc b/test/902-hello-transformation/transform.cc
deleted file mode 100644
index 3369dd4..0000000
--- a/test/902-hello-transformation/transform.cc
+++ /dev/null
@@ -1,153 +0,0 @@
-/*
- * Copyright (C) 2013 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <iostream>
-#include <pthread.h>
-#include <stdio.h>
-#include <vector>
-
-#include "art_method-inl.h"
-#include "base/logging.h"
-#include "jni.h"
-#include "openjdkjvmti/jvmti.h"
-#include "ti-agent/common_helper.h"
-#include "ti-agent/common_load.h"
-#include "utils.h"
-
-namespace art {
-namespace Test902HelloTransformation {
-
-static bool RuntimeIsJvm = false;
-
-bool IsJVM() {
-  return RuntimeIsJvm;
-}
-
-// base64 encoded class/dex file for
-//
-// class Transform {
-//   public void sayHi() {
-//     System.out.println("Goodbye");
-//   }
-// }
-const char* class_file_base64 =
-    "yv66vgAAADQAHAoABgAOCQAPABAIABEKABIAEwcAFAcAFQEABjxpbml0PgEAAygpVgEABENvZGUB"
-    "AA9MaW5lTnVtYmVyVGFibGUBAAVzYXlIaQEAClNvdXJjZUZpbGUBAA5UcmFuc2Zvcm0uamF2YQwA"
-    "BwAIBwAWDAAXABgBAAdHb29kYnllBwAZDAAaABsBAAlUcmFuc2Zvcm0BABBqYXZhL2xhbmcvT2Jq"
-    "ZWN0AQAQamF2YS9sYW5nL1N5c3RlbQEAA291dAEAFUxqYXZhL2lvL1ByaW50U3RyZWFtOwEAE2ph"
-    "dmEvaW8vUHJpbnRTdHJlYW0BAAdwcmludGxuAQAVKExqYXZhL2xhbmcvU3RyaW5nOylWACAABQAG"
-    "AAAAAAACAAAABwAIAAEACQAAAB0AAQABAAAABSq3AAGxAAAAAQAKAAAABgABAAAAEQABAAsACAAB"
-    "AAkAAAAlAAIAAQAAAAmyAAISA7YABLEAAAABAAoAAAAKAAIAAAATAAgAFAABAAwAAAACAA0=";
-
-const char* dex_file_base64 =
-    "ZGV4CjAzNQCLXSBQ5FiS3f16krSYZFF8xYZtFVp0GRXMAgAAcAAAAHhWNBIAAAAAAAAAACwCAAAO"
-    "AAAAcAAAAAYAAACoAAAAAgAAAMAAAAABAAAA2AAAAAQAAADgAAAAAQAAAAABAACsAQAAIAEAAGIB"
-    "AABqAQAAcwEAAIABAACXAQAAqwEAAL8BAADTAQAA4wEAAOYBAADqAQAA/gEAAAMCAAAMAgAAAgAA"
-    "AAMAAAAEAAAABQAAAAYAAAAIAAAACAAAAAUAAAAAAAAACQAAAAUAAABcAQAABAABAAsAAAAAAAAA"
-    "AAAAAAAAAAANAAAAAQABAAwAAAACAAAAAAAAAAAAAAAAAAAAAgAAAAAAAAAHAAAAAAAAAB4CAAAA"
-    "AAAAAQABAAEAAAATAgAABAAAAHAQAwAAAA4AAwABAAIAAAAYAgAACQAAAGIAAAAbAQEAAABuIAIA"
-    "EAAOAAAAAQAAAAMABjxpbml0PgAHR29vZGJ5ZQALTFRyYW5zZm9ybTsAFUxqYXZhL2lvL1ByaW50"
-    "U3RyZWFtOwASTGphdmEvbGFuZy9PYmplY3Q7ABJMamF2YS9sYW5nL1N0cmluZzsAEkxqYXZhL2xh"
-    "bmcvU3lzdGVtOwAOVHJhbnNmb3JtLmphdmEAAVYAAlZMABJlbWl0dGVyOiBqYWNrLTMuMzYAA291"
-    "dAAHcHJpbnRsbgAFc2F5SGkAEQAHDgATAAcOhQAAAAEBAICABKACAQG4Ag0AAAAAAAAAAQAAAAAA"
-    "AAABAAAADgAAAHAAAAACAAAABgAAAKgAAAADAAAAAgAAAMAAAAAEAAAAAQAAANgAAAAFAAAABAAA"
-    "AOAAAAAGAAAAAQAAAAABAAABIAAAAgAAACABAAABEAAAAQAAAFwBAAACIAAADgAAAGIBAAADIAAA"
-    "AgAAABMCAAAAIAAAAQAAAB4CAAAAEAAAAQAAACwCAAA=";
-
-static void JNICALL transformationHook(jvmtiEnv *jvmtienv,
-                                       JNIEnv* jni_env                 ATTRIBUTE_UNUSED,
-                                       jclass class_being_redefined    ATTRIBUTE_UNUSED,
-                                       jobject loader                  ATTRIBUTE_UNUSED,
-                                       const char* name,
-                                       jobject protection_domain       ATTRIBUTE_UNUSED,
-                                       jint class_data_len             ATTRIBUTE_UNUSED,
-                                       const unsigned char* class_data ATTRIBUTE_UNUSED,
-                                       jint* new_class_data_len,
-                                       unsigned char** new_class_data) {
-  if (strcmp("Transform", name)) {
-    return;
-  }
-  printf("modifying class '%s'\n", name);
-  bool is_jvm = IsJVM();
-  size_t decode_len = 0;
-  unsigned char* new_data;
-  std::unique_ptr<uint8_t[]> file_data(
-      DecodeBase64((is_jvm) ? class_file_base64 : dex_file_base64, &decode_len));
-  jvmtiError ret = JVMTI_ERROR_NONE;
-  if ((ret = jvmtienv->Allocate(static_cast<jlong>(decode_len), &new_data)) != JVMTI_ERROR_NONE) {
-    printf("Unable to allocate buffer!\n");
-    return;
-  }
-  memcpy(new_data, file_data.get(), decode_len);
-  *new_class_data_len = static_cast<jint>(decode_len);
-  *new_class_data = new_data;
-  return;
-}
-
-using RetransformWithHookFunction = jvmtiError (*)(jvmtiEnv*, jclass, jvmtiEventClassFileLoadHook);
-static void DoClassTransformation(jvmtiEnv* jvmtienv, JNIEnv* jnienv, jclass target) {
-  if (IsJVM()) {
-    UNUSED(jnienv);
-    jvmtienv->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_CLASS_FILE_LOAD_HOOK, nullptr);
-    jvmtiError ret = jvmtienv->RetransformClasses(1, &target);
-    if (ret != JVMTI_ERROR_NONE) {
-      char* err;
-      jvmtienv->GetErrorName(ret, &err);
-      printf("Error transforming: %s\n", err);
-    }
-  } else {
-    RetransformWithHookFunction f =
-        reinterpret_cast<RetransformWithHookFunction>(jvmtienv->functions->reserved1);
-    if (f(jvmtienv, target, transformationHook) != JVMTI_ERROR_NONE) {
-      printf("Failed to tranform class!");
-      return;
-    }
-  }
-}
-
-extern "C" JNIEXPORT void JNICALL Java_Main_doClassTransformation(JNIEnv* env,
-                                                                  jclass,
-                                                                  jclass target) {
-  JavaVM* vm;
-  if (env->GetJavaVM(&vm)) {
-    printf("Unable to get javaVM!\n");
-    return;
-  }
-  DoClassTransformation(jvmti_env, env, target);
-}
-
-// Don't do anything
-jint OnLoad(JavaVM* vm,
-            char* options,
-            void* reserved ATTRIBUTE_UNUSED) {
-  RuntimeIsJvm = (strcmp("jvm", options) == 0);
-  if (vm->GetEnv(reinterpret_cast<void**>(&jvmti_env), JVMTI_VERSION_1_0)) {
-    printf("Unable to get jvmti env!\n");
-    return 1;
-  }
-  SetAllCapabilities(jvmti_env);
-  if (IsJVM()) {
-    jvmtiEventCallbacks cbs;
-    memset(&cbs, 0, sizeof(cbs));
-    cbs.ClassFileLoadHook = transformationHook;
-    jvmti_env->SetEventCallbacks(&cbs, sizeof(jvmtiEventCallbacks));
-  }
-  return 0;
-}
-
-}  // namespace Test902HelloTransformation
-}  // namespace art
-
diff --git a/test/902-hello-transformation/transform.h b/test/902-hello-transformation/transform.h
deleted file mode 100644
index 661058d..0000000
--- a/test/902-hello-transformation/transform.h
+++ /dev/null
@@ -1,30 +0,0 @@
-/*
- * Copyright (C) 2016 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#ifndef ART_TEST_902_HELLO_TRANSFORMATION_TRANSFORM_H_
-#define ART_TEST_902_HELLO_TRANSFORMATION_TRANSFORM_H_
-
-#include <jni.h>
-
-namespace art {
-namespace Test902HelloTransformation {
-
-jint OnLoad(JavaVM* vm, char* options, void* reserved);
-
-}  // namespace Test902HelloTransformation
-}  // namespace art
-
-#endif  // ART_TEST_902_HELLO_TRANSFORMATION_TRANSFORM_H_
diff --git a/test/Android.bp b/test/Android.bp
index fe20f29..44c64c1 100644
--- a/test/Android.bp
+++ b/test/Android.bp
@@ -244,8 +244,8 @@
     defaults: ["libartagent-defaults"],
     srcs: [
         "ti-agent/common_load.cc",
+        "ti-agent/common_helper.cc",
         "901-hello-ti-agent/basics.cc",
-        "902-hello-transformation/transform.cc",
         "903-hello-tagging/tagging.cc",
         "904-object-allocation/tracking.cc",
         "905-object-free/tracking_free.cc",
diff --git a/test/ti-agent/common_helper.cc b/test/ti-agent/common_helper.cc
new file mode 100644
index 0000000..3e2b168
--- /dev/null
+++ b/test/ti-agent/common_helper.cc
@@ -0,0 +1,91 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "ti-agent/common_helper.h"
+
+#include <stdio.h>
+
+#include "jni.h"
+#include "openjdkjvmti/jvmti.h"
+#include "ti-agent/common_load.h"
+#include "utils.h"
+
+namespace art {
+bool RuntimeIsJVM;
+
+bool IsJVM() {
+  return RuntimeIsJVM;
+}
+
+void SetAllCapabilities(jvmtiEnv* env) {
+  jvmtiCapabilities caps;
+  env->GetPotentialCapabilities(&caps);
+  env->AddCapabilities(&caps);
+}
+
+namespace common_redefine {
+
+using RedefineDirectFunction = jvmtiError (*)(jvmtiEnv*, jclass, jint, const unsigned char*);
+static void DoClassTransformation(jvmtiEnv* jvmti_env, JNIEnv* env,
+                                  jclass target,
+                                  jbyteArray class_file_bytes,
+                                  jbyteArray dex_file_bytes) {
+  jbyteArray desired_array = IsJVM() ? class_file_bytes : dex_file_bytes;
+  jint len = static_cast<jint>(env->GetArrayLength(desired_array));
+  const unsigned char* redef_bytes = reinterpret_cast<const unsigned char*>(
+      env->GetByteArrayElements(desired_array, nullptr));
+  jvmtiError res;
+  if (IsJVM()) {
+    jvmtiClassDefinition def;
+    def.klass = target;
+    def.class_byte_count = static_cast<jint>(len);
+    def.class_bytes = redef_bytes;
+    res = jvmti_env->RedefineClasses(1, &def);
+  } else {
+    RedefineDirectFunction f =
+        reinterpret_cast<RedefineDirectFunction>(jvmti_env->functions->reserved3);
+    res = f(jvmti_env, target, len, redef_bytes);
+  }
+  if (res != JVMTI_ERROR_NONE) {
+    printf("Redefinition failed!");
+  }
+}
+
+// Magic JNI export that classes can use for redefining classes.
+// To use classes should declare this as a native function with signature (Ljava/lang/Class;[B[B)V
+extern "C" JNIEXPORT void JNICALL Java_Main_doCommonClassRedefinition(JNIEnv* env,
+                                                                      jclass,
+                                                                      jclass target,
+                                                                      jbyteArray class_file_bytes,
+                                                                      jbyteArray dex_file_bytes) {
+  DoClassTransformation(jvmti_env, env, target, class_file_bytes, dex_file_bytes);
+}
+
+// Don't do anything
+jint OnLoad(JavaVM* vm,
+            char* options ATTRIBUTE_UNUSED,
+            void* reserved ATTRIBUTE_UNUSED) {
+  if (vm->GetEnv(reinterpret_cast<void**>(&jvmti_env), JVMTI_VERSION_1_0)) {
+    printf("Unable to get jvmti env!\n");
+    return 1;
+  }
+  SetAllCapabilities(jvmti_env);
+  return 0;
+}
+
+}  // namespace common_redefine
+
+}  // namespace art
diff --git a/test/ti-agent/common_helper.h b/test/ti-agent/common_helper.h
index 84997f3..76543fe 100644
--- a/test/ti-agent/common_helper.h
+++ b/test/ti-agent/common_helper.h
@@ -18,9 +18,19 @@
 #define ART_TEST_TI_AGENT_COMMON_HELPER_H_
 
 #include "jni.h"
+#include "openjdkjvmti/jvmti.h"
 #include "ScopedLocalRef.h"
 
 namespace art {
+namespace common_redefine {
+
+jint OnLoad(JavaVM* vm, char* options, void* reserved);
+
+}  // namespace common_redefine
+
+extern bool RuntimeIsJVM;
+
+bool IsJVM();
 
 template <typename T>
 static jobjectArray CreateObjectArray(JNIEnv* env,
@@ -53,11 +63,7 @@
   return ret.release();
 }
 
-static void SetAllCapabilities(jvmtiEnv* env) {
-  jvmtiCapabilities caps;
-  env->GetPotentialCapabilities(&caps);
-  env->AddCapabilities(&caps);
-}
+void SetAllCapabilities(jvmtiEnv* env);
 
 }  // namespace art
 
diff --git a/test/ti-agent/common_load.cc b/test/ti-agent/common_load.cc
index a959482..2795cbc 100644
--- a/test/ti-agent/common_load.cc
+++ b/test/ti-agent/common_load.cc
@@ -23,9 +23,9 @@
 #include "base/logging.h"
 #include "base/macros.h"
 #include "common_load.h"
+#include "common_helper.h"
 
 #include "901-hello-ti-agent/basics.h"
-#include "902-hello-transformation/transform.h"
 #include "903-hello-tagging/tagging.h"
 #include "904-object-allocation/tracking.h"
 #include "905-object-free/tracking_free.h"
@@ -54,7 +54,7 @@
 // A list of all the agents we have for testing.
 AgentLib agents[] = {
   { "901-hello-ti-agent", Test901HelloTi::OnLoad, nullptr },
-  { "902-hello-transformation", Test902HelloTransformation::OnLoad, nullptr },
+  { "902-hello-transformation", common_redefine::OnLoad, nullptr },
   { "903-hello-tagging", Test903HelloTagging::OnLoad, nullptr },
   { "904-object-allocation", Test904ObjectAllocation::OnLoad, nullptr },
   { "905-object-free", Test905ObjectFree::OnLoad, nullptr },
@@ -95,6 +95,10 @@
   return true;
 }
 
+static void SetIsJVM(char* options) {
+  RuntimeIsJVM = strncmp(options, "jvm", 3) == 0;
+}
+
 extern "C" JNIEXPORT jint JNICALL Agent_OnLoad(JavaVM* vm, char* options, void* reserved) {
   char* remaining_options = nullptr;
   char* name_option = nullptr;
@@ -112,6 +116,7 @@
     printf("agent: %s does not include an OnLoad method.\n", name_option);
     return -3;
   }
+  SetIsJVM(remaining_options);
   return lib->load(vm, remaining_options, reserved);
 }
 
@@ -132,6 +137,7 @@
     printf("agent: %s does not include an OnAttach method.\n", name_option);
     return -3;
   }
+  SetIsJVM(remaining_options);
   return lib->attach(vm, remaining_options, reserved);
 }