Fault Injection with Byteman and JUnit: do even more to ensure robustness of your applications

The time when our applications lived in isolation have passed long-long ago. Nowadays applications are a very complicated beasts talking to each other using myriads of APIs and protocols, storing data in traditional or NoSQL databases, sending messages and events over the wire ...

How often did you think about what will happen if, for example, a database goes down when your application is actively querying it? Or some API endpoint suddenly starts to refuse connection? Wouldn't it be nice to have such accidents covered as part of your test suite? That's what fault injection and Byteman framework are about.

As an example, we will build a realistic, full-blown Spring application which uses Hibernate/JPA to access MySQL database and manages customers. As part of application's JUnit integration test suite, we will include three kind of test cases:

• store / find a customer
• store customer and try to query database when it's down (fault simulation)
• store customer and a database query times out (fault simulation)

There are only two preconditions for application to run on your local development box:

• MySQL server is installed and has customers database
• Oracle JDK is installed and JAVA_HOME environment variable points to it

That's being said, we are ready to go.

First, let's describe our domain model which consists from single class Customer with id and single property name. It looks as simple as that:

package com.example.spring.domain;

import java.io.Serializable;

import javax.persistence.Column;
import javax.persistence.Entity;
import javax.persistence.GeneratedValue;
import javax.persistence.Id;
import javax.persistence.Table;

@Entity
@Table( name = "customers" )
public class Customer implements Serializable{
    private static final long serialVersionUID = 1L;

    @Id
    @GeneratedValue
    @Column(name = "id", unique = true, nullable = false)
    private long id;
 
    @Column(name = "name", nullable = false)
    private String name;
  
    public Customer() {
    }
 
    public Customer( final String name ) {
        this.name = name;
    }

    public long getId() {
        return this.id;
    }

    protected void setId( final long id ) {
        this.id = id;
    }
 
    public String getName() {
        return this.name;
    }

    public void setName( final String name ) {
        this.name = name;
    }
}

For simplicity, the servicing layer is mixed with data access layer and calls database directly. Here is our CustomerService implementation:

package com.example.spring.services;

import javax.persistence.EntityManager;
import javax.persistence.PersistenceContext;

import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import com.example.spring.domain.Customer;

@Service
public class CustomerService {
    @PersistenceContext private EntityManager entityManager;
 
    @Transactional( readOnly = true )
    public Customer find( long id ) {
        return this.entityManager.find( Customer.class, id );
    }

    @Transactional( readOnly = false )
    public Customer create( final String name ) {
        final Customer customer = new Customer( name );
        this.entityManager.persist(customer);
        return customer;
    }
 
    @Transactional( readOnly = false )
    public void deleteAll() {
        this.entityManager.createQuery( "delete from Customer" ).executeUpdate();
    }
}

And lastly, the Spring application context which defines data source and transaction manager. A small note here: as we won't introduce data access layer (@Repository) classes, in order for Spring to perform exception translation properly we define PersistenceExceptionTranslationPostProcessor instance to post-process service classes (@Service). Everything else should be very familiar.

package com.example.spring.config;

import java.util.Properties;

import javax.sql.DataSource;

import org.hibernate.dialect.MySQL5InnoDBDialect;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.ComponentScan;
import org.springframework.context.annotation.Configuration;
import org.springframework.dao.annotation.PersistenceExceptionTranslationPostProcessor;
import org.springframework.jdbc.datasource.DriverManagerDataSource;
import org.springframework.orm.jpa.JpaTransactionManager;
import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean;
import org.springframework.orm.jpa.vendor.Database;
import org.springframework.orm.jpa.vendor.HibernateJpaVendorAdapter;
import org.springframework.stereotype.Service;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.annotation.EnableTransactionManagement;

import com.example.spring.services.CustomerService;

@EnableTransactionManagement
@Configuration
@ComponentScan( basePackageClasses = CustomerService.class )
public class AppConfig {
    @Bean
    public PersistenceExceptionTranslationPostProcessor exceptionTranslationPostProcessor() {
        final PersistenceExceptionTranslationPostProcessor processor = new PersistenceExceptionTranslationPostProcessor();
        processor.setRepositoryAnnotationType( Service.class );
        return processor;
    }

    @Bean
    public HibernateJpaVendorAdapter hibernateJpaVendorAdapter() {
        final HibernateJpaVendorAdapter adapter = new HibernateJpaVendorAdapter();

        adapter.setDatabase( Database.MYSQL );
        adapter.setShowSql( false );

        return adapter;
    }

    @Bean
    public LocalContainerEntityManagerFactoryBean entityManager() throws Throwable {
        final LocalContainerEntityManagerFactoryBean entityManager = new LocalContainerEntityManagerFactoryBean();
     
        entityManager.setPersistenceUnitName( "customers" );
        entityManager.setDataSource( dataSource() );
        entityManager.setJpaVendorAdapter( hibernateJpaVendorAdapter() );

        final Properties properties = new Properties();
        properties.setProperty("hibernate.dialect", MySQL5InnoDBDialect.class.getName());
        properties.setProperty("hibernate.hbm2ddl.auto", "create-drop" );
        entityManager.setJpaProperties( properties );

        return entityManager;
    }

    @Bean
    public DataSource dataSource() {
        final DriverManagerDataSource dataSource = new DriverManagerDataSource();

        dataSource.setDriverClassName( com.mysql.jdbc.Driver.class.getName() );
        dataSource.setUrl( "jdbc:mysql://localhost/customers?enableQueryTimeouts=true" );
        dataSource.setUsername( "root" );
        dataSource.setPassword( "" );

        return dataSource;
    }

    @Bean
    public PlatformTransactionManager transactionManager() throws Throwable {
        return new JpaTransactionManager( this.entityManager().getObject() );
    }
}

Now let's add a simple JUnit test case to verify our Spring application actually works as expected. Before doing that, the database customers should be created:

> mysql -u root
mysql> create database customers;
Query OK, 1 row affected (0.00 sec)

And here is a CustomerServiceTestCase which for now has single test to create the customer and verify it's actually has been created.

package com.example.spring;

import static org.hamcrest.CoreMatchers.notNullValue;
import static org.junit.Assert.assertThat;

import javax.inject.Inject;

import org.junit.After;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.test.context.ContextConfiguration;
import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
import org.springframework.test.context.support.AnnotationConfigContextLoader;

import com.example.spring.config.AppConfig;
import com.example.spring.domain.Customer;
import com.example.spring.services.CustomerService;

@RunWith( SpringJUnit4ClassRunner.class )
@ContextConfiguration(loader = AnnotationConfigContextLoader.class, classes = { AppConfig.class } )
public class CustomerServiceTestCase {
    @Inject private CustomerService customerService; 
 
    @After
    public void tearDown() {
        customerService.deleteAll();
    }

    @Test
    public void testCreateCustomerAndVerifyItHasBeenCreated() throws Exception {
        Customer customer = customerService.create( "Customer A" );
        assertThat( customerService.find( customer.getId() ), notNullValue() );
    }
}

That looks simple and straightforward. Now, let's think about scenario when customer creation succeeded but find fails because of query timeout. To do that, we need a help from Byteman.

In short, Byteman is bytecode manipulation framework. It's a Java agent implementation which runs with JVM (or attaches to it) and modifies running application bytecode as such changing its behavior. Byteman has a very good documentation and own rich set of rule definitions to perform mostly everything developer can come up with. Also, it has pretty good integration with JUnit framework. On that subject, Byteman tests are supposed to be run with @RunWith( BMUnitRunner.class ), but we already using @RunWith( SpringJUnit4ClassRunner.class ) and JUnit doesn't allow multiple test runners to be specified. Looks like a problem unless you are familiar with JUnit @Rule mechanics. It turns out that converting BMUnitRunner to JUnit rule is quite easy task:

package com.example.spring;

import org.jboss.byteman.contrib.bmunit.BMUnitRunner;
import org.junit.rules.MethodRule;
import org.junit.runners.model.FrameworkMethod;
import org.junit.runners.model.InitializationError;
import org.junit.runners.model.Statement;

public class BytemanRule extends BMUnitRunner implements MethodRule {
    public static BytemanRule create( Class< ? > klass ) {
        try {
            return new BytemanRule( klass ); 
        } catch( InitializationError ex ) { 
            throw new RuntimeException( ex ); 
        }
    }
 
    private BytemanRule( Class klass ) throws InitializationError {
        super( klass );
    }
 
    @Override
    public Statement apply( final Statement statement, final FrameworkMethod method, final Object target ) {
        Statement result = addMethodMultiRuleLoader( statement, method ); 
  
        if( result == statement ) {
            result = addMethodSingleRuleLoader( statement, method );
        }

        return result;
    }
}

And JUnit @Rule injection is as simple as that:

@Rule public BytemanRule byteman = BytemanRule.create( CustomerServiceTestCase.class );

Easy, right? The scenario we mentioned before could be rephrased a bit: when JDBC statement to select from 'customers' table is executed, we should fail with timeout exception. Here is how it looks like as JUnit test case with additional Byteman annotations:

    @Test( expected = DataAccessException.class )
    @BMRule(
        name = "introduce timeout while accessing MySQL database",
        targetClass = "com.mysql.jdbc.PreparedStatement",
        targetMethod = "executeQuery",
        targetLocation = "AT ENTRY",
        condition = "$0.originalSql.startsWith( \"select\" ) && !flagged( \"timeout\" )",
        action = "flag( \"timeout\" ); throw new com.mysql.jdbc.exceptions.MySQLTimeoutException( \"Statement timed out (simulated)\" )"
    )
    public void testCreateCustomerWhileDatabaseIsTimingOut()  {
        Customer customer = customerService.create( "Customer A" );
        customerService.find( customer.getId() );
    }

We could read it like this: "When someone calls executeQuery method of PreparedStatement class and query starts with 'SELECT' than MySQLTimeoutException will be thrown, and it should happen only once (controlled by timeout flag)". Running this test case prints stacktrace in a console and expects DataAccessException to be thrown:

com.mysql.jdbc.exceptions.MySQLTimeoutException: Statement timed out (simulated)
 at sun.reflect.NativeConstructorAccessorImpl.newInstance0(Native Method) ~[na:1.7.0_21]
 at sun.reflect.NativeConstructorAccessorImpl.newInstance(NativeConstructorAccessorImpl.java:57) ~[na:1.7.0_21]
 at sun.reflect.DelegatingConstructorAccessorImpl.newInstance(DelegatingConstructorAccessorImpl.java:45) ~[na:1.7.0_21]
 at java.lang.reflect.Constructor.newInstance(Constructor.java:525) ~[na:1.7.0_21]
 at org.jboss.byteman.rule.expression.ThrowExpression.interpret(ThrowExpression.java:231) ~[na:na]
 at org.jboss.byteman.rule.Action.interpret(Action.java:144) ~[na:na]
 at org.jboss.byteman.rule.helper.InterpretedHelper.fire(InterpretedHelper.java:169) ~[na:na]
 at org.jboss.byteman.rule.helper.InterpretedHelper.execute0(InterpretedHelper.java:137) ~[na:na]
 at org.jboss.byteman.rule.helper.InterpretedHelper.execute(InterpretedHelper.java:100) ~[na:na]
 at org.jboss.byteman.rule.Rule.execute(Rule.java:682) ~[na:na]
 at org.jboss.byteman.rule.Rule.execute(Rule.java:651) ~[na:na]
 at com.mysql.jdbc.PreparedStatement.executeQuery(PreparedStatement.java) ~[mysql-connector-java-5.1.24.jar:na]
 at org.hibernate.engine.jdbc.internal.ResultSetReturnImpl.extract(ResultSetReturnImpl.java:56) ~[hibernate-core-4.2.0.Final.jar:4.2.0.Final]
 at org.hibernate.loader.Loader.getResultSet(Loader.java:2031) [hibernate-core-4.2.0.Final.jar:4.2.0.Final]

Looks good, what about another scenario: customer creation succeeded but find fails because the database went down? This one is a bit more complicated but easy to do anyway, let's take a look:

@Test( expected = CannotCreateTransactionException.class )
@BMRules(
    rules = {
        @BMRule(
            name="create countDown for AbstractPlainSocketImpl",
            targetClass = "java.net.AbstractPlainSocketImpl",
            targetMethod = "getOutputStream",
            condition = "$0.port==3306",
            action = "createCountDown( \"connection\", 1 )"
        ),
        @BMRule(
            name = "throw IOException when trying to execute 2nd query to MySQL",
            targetClass = "java.net.AbstractPlainSocketImpl",
            targetMethod = "getOutputStream",
            condition = "$0.port==3306 && countDown( \"connection\" )",
            action = "throw new java.io.IOException( \"Connection refused (simulated)\" )"
        )
    }
)
public void testCreateCustomerAndTryToFindItWhenDatabaseIsDown() {
    Customer customer = customerService.create( "Customer A" );
    customerService.find( customer.getId() );
}

Let me explain what's going on here. We would like to sit on socket level and actually control the communication as close to network as we can, not on JDBC driver level. That's why we instrumenting AbstractPlainSocketImpl. We also know that MySQL's default port is 3306 so we are instrumenting only sockets opened on this port. Another fact, we know that first created socket corresponds to customer creation and we should let it go through. But second one corresponds to find and must fail. The createCountDown named "connection" serves this purposes: the first call goes through (latch doesn't count to zero yet) but second call triggers MySQLTimeoutException exception. Running this test case prints stacktrace in a console and expects CannotCreateTransactionException to be thrown:

Caused by: java.io.IOException: Connection refused (simulated)
 at sun.reflect.NativeConstructorAccessorImpl.newInstance0(Native Method) ~[na:1.7.0_21]
 at sun.reflect.NativeConstructorAccessorImpl.newInstance(NativeConstructorAccessorImpl.java:57) ~[na:1.7.0_21]
 at sun.reflect.DelegatingConstructorAccessorImpl.newInstance(DelegatingConstructorAccessorImpl.java:45) ~[na:1.7.0_21]
 at java.lang.reflect.Constructor.newInstance(Constructor.java:525) ~[na:1.7.0_21]
 at org.jboss.byteman.rule.expression.ThrowExpression.interpret(ThrowExpression.java:231) ~[na:na]
 at org.jboss.byteman.rule.Action.interpret(Action.java:144) ~[na:na]
 at org.jboss.byteman.rule.helper.InterpretedHelper.fire(InterpretedHelper.java:169) ~[na:na]
 at org.jboss.byteman.rule.helper.InterpretedHelper.execute0(InterpretedHelper.java:137) ~[na:na]
 at org.jboss.byteman.rule.helper.InterpretedHelper.execute(InterpretedHelper.java:100) ~[na:na]
 at org.jboss.byteman.rule.Rule.execute(Rule.java:682) ~[na:na]
 at org.jboss.byteman.rule.Rule.execute(Rule.java:651) ~[na:na]
 at java.net.AbstractPlainSocketImpl.getOutputStream(AbstractPlainSocketImpl.java) ~[na:1.7.0_21]
 at java.net.PlainSocketImpl.getOutputStream(PlainSocketImpl.java:214) ~[na:1.7.0_21]
 at java.net.Socket$3.run(Socket.java:915) ~[na:1.7.0_21]
 at java.net.Socket$3.run(Socket.java:913) ~[na:1.7.0_21]
 at java.security.AccessController.doPrivileged(Native Method) ~[na:1.7.0_21]
 at java.net.Socket.getOutputStream(Socket.java:912) ~[na:1.7.0_21]
 at com.mysql.jdbc.MysqlIO.(MysqlIO.java:330) ~[mysql-connector-java-5.1.24.jar:na]

Great! The possibilities Byteman provides for different fault simulations are enormous. Carefully adding test suites to verify how application reacts to erroneous conditions greatly improves application robustness and resiliency to failures. Bunch of thanks to Byteman guys!

Please find the complete project on GitHub.