Mock Testing: Write a Java Agent Plugin to Mock Your Custom Components

• Mock Testing with Java agent
• AREX Agent
• Setting up the development environment
• Step 1: Create DalClientModuleInstrumentation Class
• Step 2: Bytecode Modification
• Step 3: Implementation of recording and replay
• Deployments
• Debug

Mock Testing with Java agent

Mock testing is a software testing method used to simulate the behavior of certain components or dependencies in a system. The purpose of mock testing is to isolate the system under test and verify the correctness of interactions between the system and its dependencies.

Java Agent is a Java technology that allows application startup modification and enhancement of bytecode. With Java Agent, it is possible to dynamically modify the behavior of target classes at runtime, including intercepting method calls and modifying return values. This provides a possibility for implementing Mock.

By implementing a proxy class or using bytecode manipulation libraries like Byte Buddy or ASM, you can modify the behavior of the target classes. This allows you to return mock data or execute custom logic instead of the actual behavior of the dependencies.

By loading Java Agent into the target application, it can intercept and modify calls to open source components at application startup, thus enabling Mock of open source components. This can isolate the system under test during testing and verify the correctness of interactions between the system and its dependencies.

AREX Agent

AREX provides an out-of-box agent file that could be attached to any applications with Java 8+ and dynamically weaves solid bytecode into your existing code to record the real data of live traffic, and further use and replay it for mocking, testing, and debugging purposes.

Let's take a simple component DalClient as an example to illustrate how AREX Agent implements mocking.

Object result = dalClient.query(key)
if (result != null) {
    return result; // return the value when recording
} else {
    return rpc.search(); // no value to return thus call the remote interface when replaying

Setting up the development environment

1. Download the code and run mvn install.

git clone https://github.com/arextest/arex-agent-java.git 
mvn clean install // install arex-agent-java dependencies that may be needed into your local maven repository.

1. Create a new Java project from IntelliJ IDEA.

Now, add the AREX Agent dependencies to the pom file:

<dependencies>
    <!-- Some basic classes and tools used by arex-agent, such as recording and playback related components, configuration information, etc. -->         
    <dependency>
        <groupId>io.arex</groupId>
        <artifactId>arex-instrumentation-api</artifactId>
        <version>${arex.version}</version>
    </dependency>
    <!-- serialization related components -->         
    <dependency>
        <groupId>io.arex</groupId>
        <artifactId>arex-serializer</artifactId>
        <version>${arex.version}</version>
    </dependency>
    <!-- Replace it with the component you need to mock in the actual development. -->
    <dependency>
        <groupId>com.your.company.dal</groupId>
        <artifactId>dalclient</artifactId>
        <version>1.0.0</version>
        <scope>provided</scope>
    </dependency>
</dependencies>

Step 1: Create DalClientModuleInstrumentation Class

The arex-agent-java loads and instantiates plugins using the SPI (Service Provider Interface) mechanism. Therefore, we declare a class called DalClientModuleInstrumentation with the com.google.auto.service.AutoService annotation. This class serves as the entry point for modifying DalClient and will be recognized by the arex-agent-java (@AutoService).

@AutoService(ModuleInstrumentation.class)
public class DalClientModuleInstrumentation extends ModuleInstrumentation {
    public DalClientModuleInstrumentation() {
       // Plugin module name, you can specify a different version match if the code difference between different versions of your DalClient component is relatively large and to be supported in separate versions:
       // ModuleDescription.builder().name("dalclient").supportFrom(ComparableVersion.of("1.0")).supportTo(ComparableVersion.of("2.0")).build();
       super("plugin-dal");
    }

    @Override
    public List<TypeInstrumentation> instrumentationTypes() {
      // classes that have undergone bytecode modification to change the behavior of the DalClient component
      return singletonList(new DalClientInstrumentation()); 
    }
}

So when it's necessary to distinguish between version numbers?

If there're differences between two versions of the DalClient component, eg., the name of "invoke()" method in version 1.0.0 is changed to "invokeAll()" in version 2.0.0. In this case, the code modified by the Agent plugin cannot cover both versions of the DalClient framework simultaneously. Therefore, adapting the code specifically for different versions may be necessary.

The implementation details can refer to the adaptation logic for different Dubbo versions in the arex-instrumentation/dubbo/ module of the arex-agent-java project.

The version matching is implemented based on the contents of the META-INF/MANIFEST.MF file in the jar package of the component.

Step 2: Bytecode Modification

To modify the bytecode, create a new file named DalClientInstrumentation.java.

Modification of the DalClient source code is quite simple. Just find the underlying implementation methods, preferably using a common API, and then using bytecode manipulation tools like Byte Buddy to modify this API by adding our own code to implement the mocking functionality.

Here's the DalClient source code to be modified:

package com.your.company.dal;

public class DalClient {
    public Object query(String param) {
        return this.invoke(DalClient.Action.QUERY, param);
    }

    public Object insert(String param) {
        return this.invoke(DalClient.Action.INSERT, param);
    }

    private Object invoke(Action action, String param) {
      Object result;         
        switch (action) {
            case QUERY:
                result = "query:" + param;
            case INSERT:
                result = "insert:" + param;
            case UPDATE:
                result = "update:" + param;
            default:
                result = "unknown action:" + param;
        }     
        return result;
    }

    public static enum Action {
        QUERY,
        INSERT,
        UPDATE;
        private Action() {
        }
    }
}

In regular business projects, the DalClient is typically invoked using the dalClient.query(key) method. From the above source code, it can be observed that the underlying implementation is achieved through the invoke method.

Therefore, by modifying the invoke method and injecting recording and playback code, the functionality of the DalClientInstrumentation class can be summarized as follows:

public class DalClientInstrumentation extends TypeInstrumentation {
    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {         
        // Path to the DalClient class to be modified
        return named("com.your.company.dal.DalClient"); 
    }

    @Override
    public List<MethodInstrumentation> methodAdvices() {
        ElementMatcher<MethodDescription> matcher = named("invoke") // the method tobe modified
                .and(takesArgument(0, named("com.your.company.dal.DalClient$Action"))) // The first parameter type of this method. This is to distinguish that there may be methods with the same name but different parameter types
                .and(takesArgument(1, named("java.lang.String"))); 
            // The InvokeAdvice class is the code we need to inject in the invoke method         
            return singletonList(new MethodInstrumentation(matcher, InvokeAdvice.class.getName())); 
    }
}

The MethodInstrumentation, ElementMatcher, named, takesArgument, and other methods in the above code are all the ByteBuddy API. AREX Agent edits bytecode with ByteBuddy for recording and replay. (See more details, refer to: https://bytebuddy.net/#/tutorial)

Step 3: Implementation of recording and replay

Here is the code injected into the DalClient#invoke method, implementing InvokeAdvice :

public static class InvokeAdvice {
    // OnMethodEnter is the operation that executes before the modified method (invoke) logic is called
    @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class, suppress = Throwable.class)
    public static boolean onEnter(@Advice.Argument(0) DalClient.Action action, // Obtaining a reference to the first parameter (action) of the modified method
                                  @Advice.Argument(1) String param, // Obtaining a reference to the second parameter (param) of the modified metho
                                  @Advice.Local("mockResult") MockResult mockResult) { // mockResult is the customized variables within this method 
        mockResult = DalClientAdvice.replay(action, param); // replay
        return mockResult != null && mockResult.notIgnoreMockResult();
    }

    // OnMethodExit is the operation that executes before the end of the modified method (invoke)
    @Advice.OnMethodExit(suppress = Throwable.class)   
    public static void onExit(@Advice.Argument(0) DalClient.Action action,
                              @Advice.Argument(1) String param,
                              @Advice.Local("mockResult") MockResult mockResult,
                              @Advice.Return(readOnly = false) Object result) { 
        // returned result of the method
        if (mockResult != null && mockResult.notIgnoreMockResult()) {
            result = mockResult.getResult(); // assigns the replayed result to the result variable
            return;
        }
        DalClientAdvice.record(action, param, result); // recording logic
    }
}

This class is designed to inject code before and after the modified invoke method is invoked, enabling the implementation of recording and replay functionality.

In which:

The skipOn = Advice.OnNonDefaultValue parameter indicates that if mockResult != null && mockResult.notIgnoreMockResult() is true (non-default value, where the default value for a boolean type is false), the original logic of the method will be skipped. In other words, the original method logic will not be executed, and instead, the mocked value will be returned. If it is false, the original method logic will be executed as usual.

The modified bytecode is shown below:

public class DalClientInstrumentation extends TypeInstrumentation {
    private Object invoke(DalClient.Action action, String param) {
        // replay          
        MockResult mockResult = DalClientAdvice.replay(action, param);
        if (mockResult != null && mockResult.notIgnoreMockResult()) {
            return mockResult.getResult();
        }

        // the original logic
        Object result;
        switch (action) {
            case QUERY:
                result = "query:" + param;
            case INSERT:
                result = "insert:" + param;
            case UPDATE:
                result = "update:" + param;
            default:
                result = "unknown action:" + param;
        }

        DalClientAdvice.record(action, param, result); // record

        return result;
    }
}

Similar to the Aspect-Oriented Programming (AOP), this class inserts our code before and after the method invocation. In the case of replay, it returns the mocked result without executing the original logic. For recording, it records the result before it returns.

The code of DalClientAdvice is shown below:

public class DalClientAdvice {
    // record
    public static void record(DalClient.Action action, String param, Object result) {
        if (ContextManager.needRecord()) {
            Mocker mocker = buildMocker(action, param);
            mocker.getTargetResponse().setBody(Serializer.serialize(result));
            MockUtils.recordMocker(mocker);
        }
    }

    // replay
    public static MockResult replay(DalClient.Action action, String param) {
        if (ContextManager.needReplay()) {
            Mocker mocker = buildMocker(action, param);
            Object result = MockUtils.replayBody(mocker);
            return MockResult.success(result); DalClientInstrumentation          }
        return null;
    }

    private static Mocker buildMocker(DalClient.Action action, String param) {
        Mocker mocker = MockUtils.createDatabase(action.name().toLowerCase());
        mocker.getTargetRequest().setBody(param);
        return mocker;
    }
}

The above is just a simple demo. To see more details, refer to the arex-instrumentation module of the arex-agent-java project for specific usage, which has a lot of implementations that modify various middleware.

Summarized in three steps:

1. Create a entry class called DalClientModuleInstrumentation, which will be loaded when arex-agent-java starts.
2. Create a class called DalClientInstrumentation to specify which class and method to modify, as well as when to execute the modifications, and inform the AREX Agent accordingly.
3. Create a class called DalClientAdvice to implement the actual recording and replay functionality, or any custom logic you want to implement.

Deployments

After completing the development, you can test whether your plugin works correctly by following these steps:

1. Run the command mvn clean compile package to generate the plugin JAR file (it will be located in the /target directory of your project by default).
2. In the arex-agent-java project, run the command mvn clean compile package to generate the arex-agent-jar directory (located in the root directory of the project).
3. Create a new folder named extensions in the arex-agent-jar directory to store the extension JAR files.
4. Place the generated plugin JAR file (e.g., plugin-dal-1.0.0.jar) into the extensions folder.

The AREX Agent loads all JAR files in the extensions folder as extension features during startup. The directory structure is as follows:

Now you can debug your plugin. Start your business application, making sure to attach the agent in the VM options:

The detailed parameters are as follows:

-javaagent:<your-project>\arex-agent-java\arex-agent-0.3.8.jar
-Darex.service.name=<your-service-name>
-Darex.storage.service.host=<arex-storage-service-ip:port>
-Darex.enable.debug=true

By starting the project in this way, the agent will be able to load the plugin package.

If the project references the internal DalClient component, you will be able to see the name of the plugin declared in DalClientModuleInstrumentation in the console after starting the project.

If the console outputs the log message [arex] installed instrumentation module: plugin-dal, it means that the plugin has been recognized and loaded successfully.

Next, test if the plugin is functioning correctly. As shown below, if an interface in the project invokes the query method of the DalClient component, you can observe whether the plugin can record or replay the invoke method of the DalClient component. If the recording is successful, the following log message will be printed:

Similarly, to test the replay functionality, include the header arex-record-id:AREX-10-32-179-120-2126724921 (the recordId generated during recording) in the request header. After the successful replay, the console will also output log messages related to [arex] replay category: Database, operation: query.

Debug

If the plugin is not working properly or not behaving as expected, you can troubleshoot from the following two points:

1. Check the bytecode in the arex-agent-jar/bytecode-dump folder。

The bytecode-dump folder holds the modified bytecode files where you can find the class files before and after modification here:

Drag and drop the above two files into your IDE to open them directly. Confirm whether the modifications were successful. You can compare the code logic before (left side) and after (right side) the plugin modifications, as shown in the comparison image below:

• If the modifications were not successful, please refer to the source code of the project demo at GitHub.

• If the code modifications were successful but the behavior is not as expected, you can debug the plugin source code in your IDE by following these steps:

1. Import the arex-agent-java project into the IDE of your project by following these steps: File → New → Module from Existing Sources

Select the local arex-agent-java project, and then choose Maven to import it.

1. Import the plugin project into the IDE of your project. Then you will have the source code for the business project, as well as the arex-agent-java and plugin-dal projects.

This way, you can set breakpoints in both of these projects to debug the code of the Agent and the plugin, as shown in the following image:

Community⤵️
🐦 Follow us on Twitter
📝 Join AREX Slack
📧 Join the Mailing List