diff --git a/test/904-object-allocation/tracking.cc b/test/904-object-allocation/tracking.cc
index c392cd4..b22fc6c 100644
--- a/test/904-object-allocation/tracking.cc
+++ b/test/904-object-allocation/tracking.cc
@@ -26,13 +26,12 @@
 #include "openjdkjvmti/jvmti.h"
 #include "ScopedLocalRef.h"
 #include "ScopedUtfChars.h"
+#include "ti-agent/common_load.h"
 #include "utils.h"
 
 namespace art {
 namespace Test904ObjectAllocation {
 
-static jvmtiEnv* jvmti_env;
-
 static std::string GetClassName(JNIEnv* jni_env, jclass cls) {
   ScopedLocalRef<jclass> class_class(jni_env, jni_env->GetObjectClass(cls));
   jmethodID mid = jni_env->GetMethodID(class_class.get(), "getName", "()Ljava/lang/String;");
@@ -58,7 +57,8 @@
          static_cast<size_t>(size));
 }
 
-extern "C" JNIEXPORT void JNICALL Java_Main_setupCallback(JNIEnv* env ATTRIBUTE_UNUSED, jclass) {
+extern "C" JNIEXPORT void JNICALL Java_Main_setupObjectAllocCallback(
+    JNIEnv* env ATTRIBUTE_UNUSED, jclass klass ATTRIBUTE_UNUSED) {
   jvmtiEventCallbacks callbacks;
   memset(&callbacks, 0, sizeof(jvmtiEventCallbacks));
   callbacks.VMObjectAlloc = ObjectAllocated;
@@ -82,7 +82,7 @@
   if (ret != JVMTI_ERROR_NONE) {
     char* err;
     jvmti_env->GetErrorName(ret, &err);
-    printf("Error getting tag: %s\n", err);
+    printf("Error enabling/disabling allocation tracking: %s\n", err);
   }
 }