EJB, Dependency Injection and Unit Testing

Our team spent the last week in Paspels to improve our programming and testing skills. During that week we played around with OpenEJB to test EJB’s in an embedded JEE container. To get to know the OpenEJB framework we implemented some tests. While integration tests worked just fine, writing real “FIRST” unit tests posed some challenges. It is not possible to mock related (injected) objects and the startup of the container is time-consuming.

Our first approach to mock the injected EJB’s was an additional constructor to initialize them. This solution is ugly because we must not extend production code just for unit testing. Therefore, we decided to implement our own simple EJB injector.

As an example we developed a blackboard application. The server part is implemented using stateless session beans. One of them is the blackboard service. Besides other features, it provides a method for deleting messages.

@Stateless
public class BlackboardServiceImpl implements BlackboardService {

  @PersistenceContext(unitName = "pinacolada-unit")
  private EntityManager entityManager;

  @EJB
  private SessionTokenManager sessionTokenManager;

  @EJB
  private TimeServiceLocal timeService;

  @Override
  public void deleteMessage(final Message msg, final String sessionToken) throws NoSessionException {
    sessionTokenManager.validateToken(sessionToken);
    entityManager.remove(msg);
  }
  // implementation...
}

To setup the test environment of our BlackboardService we wanted to use code as listed below. The first lines create the mocks for the injected EJBs using Mockito. Next we create the injector, assign the mappings and pass over the bean to initialize its annotated fields.

@Before
public void test() throws Exception {
  this.testee = new BlackboardServiceImpl();
  // create mocks
  this.entityManagerMock = mock(EntityManager.class);
  this.sessionTokenManagerMock = mock(SessionTokenManager.class);
  this.timeServiceMock = mock(TimeServiceLocal.class);
  // inject
  final EJBTestInjector injector = new EJBTestInjector();
  injector.assign(EntityManager.class, entityManagerMock);
  injector.assign(SessionTokenManager.class, sessionTokenManagerMock);
  injector.assign(TimeServiceLocal.class, timeServiceMock);
  injector.inject(this.testee);
}

With this setup we are able to test our blackboard service fully isolated. We don’t need to start up an embedded EJB container which makes our unit tests fast.

The version of the injector we published below was written in a short pair programming session. We are aware of the simplicity of our implementation. It does not support the more sophisticated features of the EJB specification. But it does a good job for simple EJBs as we see them every day.

import java.lang.annotation.Annotation;
import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import javax.annotation.Resource;
import javax.ejb.EJB;
import javax.persistence.PersistenceContext;

public class EJBTestInjector {

  private static final List<Class<? extends Annotation>> EJB_ANNOTATIONS;
  static {
    EJB_ANNOTATIONS = new ArrayList<Class<? extends Annotation>>();
    EJB_ANNOTATIONS.add(EJB.class);
    EJB_ANNOTATIONS.add(PersistenceContext.class);
    EJB_ANNOTATIONS.add(Resource.class);
  }

  final Map<Class<?>, Object> mappings = new HashMap<Class<?>, Object>();

  public void inject(final Object bean) throws Exception {
    for (final Field field : getEJBAnnotatedFields(bean)) {
       injectField(field, bean);
    }
  }

  public void assign(final Class<?> type, final Object instance) {
    mappings.put(type, instance);
  }

  private void injectField(final Field field, final Object bean)  throws Exception {
    final Object instanceToInject = mappings.get(field.getType());
    if (!field.isAccessible()) {
      field.setAccessible(true);
    }
    field.set(bean, instanceToInject);
  }

  private List<Field> getEJBAnnotatedFields(final Object bean) {
    final Class<? extends Object> beanClass = bean.getClass();
    final List<Field> annotatedFields = new ArrayList<Field>();
    for (final Field field : beanClass.getDeclaredFields()) {
      if (hasEJBAnnotation(field)) {
        annotatedFields.add(field);
      }
    }
    return annotatedFields;
  }

  private static boolean hasEJBAnnotation(final Field field) {
    for (final Class<? extends Annotation> annotation : EJB_ANNOTATIONS) {
      if (field.isAnnotationPresent(annotation)) {
        return true;
      }
    }
    return false;
  }
}

Our injector iterates over all EJB annotated fields (EJB, PersistenceContext, Resource) and tries to inject them with one of the specified instances. That’s it. We are amazed of the simplicity of the injector and the benefit it brings to the testability in the realm of JEE.