Unit Tests with Guice and Warp Persist

Yesterday I needed to do some unit testing for a project using Guice and Warp Persist. And one thing that seems to be lacking, or wasn’t to be found be me yesterday, is the kind of unit testing you can do with Spring, where each method is wrapped by a transaction which is rolled back when the method ends.

To simplify, it enables you to do this:

@Test
public void createUser() {
    User u = userService.createUserWithId(1);
}

@Test
public void createUserWithSameId() {
    User u = userService.createUserWithId(1);
}

If you assume the “userService” is transactional, we’re on a database that spans multiple methods, and that the user ID must be unique, the above pseudo-code should fail as we’re trying to create two users with the same ID. If, however, there was a transaction wrapping both methods, and that transaction was rolled back, we’d be fine.

So, can you do it with Guice and Warp Persist? Sure you can!

(I’m using TestNG and JPA by the way, but you’ll get the point).

We’ll create a base class that sets up the Guice context as well as handles the transactions for us. Let’s start with creating the methods we need:

public abstract class JpaTestBase {

  @BeforeClass
  public void setUpJpa() throws Exception {
    // setup guice + jpa here here
  }

  @AfterClass
  public void tearDownJpa() throws Exception {
    // stop jpa here
  }

  /**
   * Return the module needed for the test.
   */
  protected abstract List<? extends Module> testModules();

  @BeforeMethod
  public void setupTransaction() throws Exception {
    // create "session in view"
  }

  @AfterMethod
  public void cleanTransaction() throws Exception {
    // rollback transaction and close session
  }
}

We’re using the TestNG “before” and “after” class to setup and tear down JPA. If the test was in a bigger suite you could probably do it around all test classes, but this will do for now. I’m also including a “testModules” method that subclasses should use to make sure their own test classes are setup correctly.

Before we go on, I’ll add some dependencies which will be explained later:

@Inject
protected PersistenceService service;

@Inject
protected WorkManager workManager;

@Inject
protected Provider<EntityManagerFactory> emfProvider;

We’ll setup the Guice context and the JPA persistence service in the “setUpJpa” method:

@BeforeClass
public void setUpJpa() throws Exception {
  // create list with subclass modules
  List<Module> list = new LinkedList<Module>(testModules());

  // add persistence module
  list.add(PersistenceService
            .usingJpa()
            .across(UnitOfWork.REQUEST)
            .forAll(Matchers.annotatedWith(Transactional.class), Matchers.any())
            .buildModule());

  // modules to array and create Guice
  Module[] arr = list.toArray(new Module[list.size()]);
  injector = Guice.createInjector(arr);

  // make sure we get our dependencies
  injector.injectMembers(this);

  // NB: we need to start the service (injected)
  service.start();
}

The persistence service will work across UnitOfWork.REQUEST as that’s what we’re emulating here. I’m also matching the transaction for any classes, this enables me to mark an entire class as transactional, as opposed to single methods, which I find handy.

All we need to do to shut down is to close the service, like so:

@AfterClass
public void tearDownJpa() throws Exception {
  service.shutdown();
}

Now, let’s wrap the methods. Remember our dependency injections earlier? Well, here’s where they come in handy. The WorkManager is used by Warp Persist to manage a session, we can tell it to start and end “work” and this will correspond to an EntityManager bound to the current thread. Let’s start with that:

@BeforeMethod
public void setupTransaction() throws Exception {
  workManager.beginWork();
}

@AfterMethod
public void cleanTransaction() throws Exception {
  workManager.endWork();
}

So before each method we’ll open a new EntityManager which will be closed when the method ends. So far so good. Now, in order to enable a rollback when the method ends we first need to wrap the entire execution in a transaction, which means we need to get hold of the EntityManager. Luckily, Warp Persist has a class called ManagedContext which holds currently bound objects (a bit like a custom scope), in which we’ll find our EntityManager keyed to it’s persistence context, ie. the EntityManagerFactory that created it. Take a look again at the injected dependencies we injected above: As the test only handles one persistence context we can safely let Guice inject a Provider for an EntityManagerFactory for us and assume it is going to be the right one.

Still with me? OK, let’s do it then:

@BeforeMethod
public void setupTransaction() throws Exception {
  // begin work
  workManager.beginWork();
  // get the entity manager, using it's factory
  EntityManagerFactory emf = emfProvider.get();
  EntityManager man = ManagedContext.getBind(EntityManager.class, emf);
  // begin transaction
  man.getTransaction().begin();
}

@AfterMethod
public void cleanTransaction() throws Exception {
  // get the entity manager, using its factory
  EntityManagerFactory emf = emfProvider.get();
  EntityManager man = ManagedContext.getBind(EntityManager.class, emf);
  // rollback transaction
  man.getTransaction().rollback();
  // end work
  workManager.endWork();
}

And that’s it! Each test method is now within a transaction that will be rolled back when the method ends. Now all you need to do is to write the actual tests…