How to mock a method call with Java Agent?
Tags: devops, java, testing, opensource
Published: 2023-11-29 06:00:59

<!-- TOC start (generated with https://github.com/derlin/bitdowntoc) -->

- [Code example](#code-example)
- [The implementation in AREX](#the-implementation-in-arex)
- [Record and replay in-memory data](#record-and-replay-in-memory-data)

<!-- TOC end -->

AREX is an open-source automated regression testing platform based on real requests and data. It enables the accomplishment of bytecode injection with Java Agent to record real traffic in the production environment, and then replays the requests and injects mock data in the testing environment. This enables automatic recording, replaying, and comparison, providing convenience for interface regression testing.

In general, the basic principles of implementing mocking using Java Agent can be summarized as follows:

1. Start the application with an Java Agent loaded.
2. In the premain method of the Agent, use a Java bytecode manipulation library like ByteBuddy to modify the bytecode of the target class.
3. Locate the method or field in the target class where mocking needs to be applied within the bytecode.
4. Use the bytecode manipulation library to modify the implementation of the target method or field to incorporate the mocking logic. This can involve replacing the original method with a mock method or modifying the value of a field.
5. Redefine the modified bytecode as a new class and load it into the JVM.
6. During runtime, the modified class will be used, enabling the mocking functionality.

AREX Java Agent supports mocking for not only various technology frameworks, but also local time, cached data, and various memory data types. This allows for precise reproduction of the data environment during production execution during replay, without generating dirty data to your database.

This blog will delve into how AREX achieves automatic mocking of data during traffic replay from a code perspective.

### Code example

Let's take a simple function as an example to understand the implementation. It's a function that converts a given IP string to an integer value. The code is as follows:

```java
public Integer parseIp(String ip) {
 	int result = 0;
 	if (checkFormat(ip)) { // Check if the IP string is legal
 		String[] ipArray = ip.split("\\.");
 		for (int i = 0; i < ipArray.length; i++) {
 			result = result << 8;
 			result += Integer.parseInt(ipArray[i]);
 		}
 	}
 	return result;
}
```

Now, let's explain from two perspectives, how to implement traffic playback with the function mock:

- Record (collect traffic)

When this function is called, we collect the request parameters and corresponding results for later use in traffic replay. The code is as follows:

```java
if (needRecord()) {
	// Data collection, save the parameters and results in DB
	DataService.save("parseIp", ip, result);
}
```

- Replay

When traffic replay, the previously collected data is used to implement mocking of this function. The code is as follows:

```java
if (needReplay()) {
	return DataService.query("parseIp", ip);
}
```

By examining the complete code, we can better understand its logic:

```java
public Integer parseIp(String ip) {
	if (needReplay()) {
		// in the case of replay, the collected data is used as the return result, i.e., Mock
		return DataService.query("parseIp", ip);
	}

	int result = 0;
	if (checkFormat(ip)) {
		String[] ipArray = ip.split("\\.");
		for (int i = 0; i < ipArray.length; i++) {
			result = result << 8;
			result += Integer.parseInt(ipArray[i]);
		}
	}

	if (needRecord()) {
		// in the case of recording, save the parameters and results in DB
		DataService.save("pareseIp", ip, result);
	}
	return result;
}
```

### The implementation in AREX

The implementation in AREX is similar but more complex. The developers don't need to add recording and replaying code in the business logic. The `arex-agent` automatically adds the necessary code in the required code. Let's take the Query of MyBatis3 as an example.

By reading the MyBatis source code, we should know that all Query operations are consolidated in the `query` method of the `org.apache.ibatis.executor.BaseExecutor` class (excluding Batch operations). The method signature is as follows:

```java
public <E> List<E> query(MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, CacheKey key, BoundSql boundSql) throws SQLException
```

This query method includes the SQL query and accompanying parameters. The function's result includes the data read from the database. It is obviously appropriate to collect data here, and during replay, the collected data can be used as the result to avoid actual operations performed on database. 

Let's take a look at the code in AREX. For the sake of simplicity, some simplifications have been made, as shown below:

```java
public class ExecutorInstrumentation extends TypeInstrumentation {
    @Override
    protected ElementMatcher<TypeDescription> typeMatcher() {
        // Full name of the class which need code injection
        return named("org.apache.ibatis.executor.BaseExecutor");
    }

    @Override
    public List<MethodInstrumentation> methodAdvices() {
        // The name of the method that requires code injection. Because the query method has multiple overloads, parameter verification is included.
        return Collections.singletonList(new MethodInstrumentation(
                        named("query").and(isPublic())
                                .and(takesArguments(6))
                                .and(takesArgument(0, named("org.apache.ibatis.mapping.MappedStatement")))
                                .and(takesArgument(1, Object.class))
                                .and(takesArgument(5, named("org.apache.ibatis.mapping.BoundSql"))),
                        QueryAdvice.class.getName())
        );
    }

    // injected code
    public static class QueryAdvice {

        @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class, suppress = Throwable.class)
        public static boolean onMethodEnter(@Advice.Argument(0) MappedStatement var1,
                                            @Advice.Argument(1) Object var2,
                                            @Advice.Argument(5) BoundSql boundSql,
                                            @Advice.Local("mockResult") MockResult mockResult) {
            RepeatedCollectManager.enter(); // prevent repeated collection of data caused by nested calls
            if (ContextManager.needReplay()) {
                mockResult = InternalExecutor.replay(var1, var2, boundSql, "query");
            }
            return mockResult != null;
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onMethodExit(@Advice.Argument(0) MappedStatement var1,
                                  @Advice.Argument(1) Object var2,
                                  @Advice.Argument(5) BoundSql boundSql,
                                  @Advice.Thrown(readOnly = false) Throwable throwable,
                                  @Advice.Return(readOnly = false) List<?> result,
                                  @Advice.Local("mockResult") MockResult mockResult) {
            if (!RepeatedCollectManager.exitAndValidate()) {
                return;
            }

            if (mockResult != null) {
                if (mockResult.getThrowable() != null) {
                    throwable = mockResult.getThrowable();
                } else {
                    result = (List<?>) mockResult.getResult();
                }
                return;
            }            

            if (ContextManager.needRecord()) {
                InternalExecutor.record(var1, var2, boundSql, result, throwable, "query");
            }
        }
    }
}
```

Where `QueryAdvice` is the code that needs to be injected in the `query` method. The code injected through `onMethodEnter` will be executed at the beginning of the method, while the code injected by `onMethodExit` will be executed before the function returns the result.

It may be difficult to understand, so let's dump the code of BaseExecutor's query method after code injection and analyze it.

```java
public <E> List<E> query(MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, CacheKey key, BoundSql boundSql) throws SQLException {
        MockResult mockResult = null;
        boolean skipOk;
        try {
            RepeatedCollectManager.enter();
            if (ContextManager.needReplay()) {
                mockResult = InternalExecutor.replay(ms, parameter, boundSql, "query");
            }
 
            skipOk = mockResult != null;
        } catch (Throwable var28) {
            var28.printStackTrace();
            skipOk = false;
        }
 
        List result;
        Throwable throwable;
        if (skipOk) {
            // in the case of replay, the original query method body is no longer executed
            result = null;
        } else {
            try {
                // The source code of the query method in BaseExecutor(omitted here), and the only adjustment is to modify the code that is returned in the original method to assign the result to a variable named result
                result = list;
            } catch (Throwable var27) {
                throwable = var27;
                result = null;
            }
        }
 
        try {
            if (mockResult != null) {
                if (mockResult.getThrowable() != null) {
                    throwable = mockResult.getThrowable();
                } else {
                    result = (List)mockResult.getResult();
                }
            } else if (RepeatedCollectManager.exitAndValidate() && ContextManager.needRecord()) {
                InternalExecutor.record(ms, parameter, boundSql, result, throwable, "query");
            }
        } catch (Throwable var26) {
            var26.printStackTrace();
        }
 
        if (throwable != null) {
            throw throwable;
        } else {
            return result;
        }
    }
```

As you see, the code of `onMethodEnter` and `onExit` is inserted at the beginning and end respectively. 

- Record

When send a request to a service, AREX will decide whether to record this request based on the configured recording frequency. If recording is required, all 3rd dependencies accessed during this request will be recorded.

AREX will record this request by invoking the `InternalExecutor.record(ms, parameter, boundSql, result, throwable, "query")` method to store the returned result, core parameters, and other information into AREX's database.

- Replay

When resend the recorded request, the function will not be executed actually. Instead, AREX will directly return the recorded result previously by invoking `InternalExecutor.replay(ms, parameter, boundSql, "query")`.

### Record and replay in-memory data

The previous example is idempotent function. For idempotent functions, since the returned result remains the same regardless of external factors, there is no need for collecting data and mocking during the recording and replay process.

For non-idempotent functions, the results returned by the request are different each time due to the influence of the external environment, such as local caches where different environment data may lead to different output results. In such cases, AREX can also record and mock the non-idempotent functions by configure it as dynamic classes.

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/gjav4jclnawo0vj5lldt.png)

Here, you can configure the class name, method name (optional, if not configured, it will be applied to all public methods with parameters and return values), and parameter types (optional) Once configured, the `arex-agent` will automatically inject similar code into the corresponding methods, enabling data collection and mocking during the replay process.

---

Community⤵️
🐦 Follow us on [Twitter](https://twitter.com/AREX_Test)
📝 Join AREX [Slack](https://arexcommunity.slack.com/ssb/redirect)
📧 Join the [Mailing List](https://groups.google.com/g/arex-test)