Introduction

In JPA when we want to insert an entity in the database we use the EntityManager.persist(..) method. In addition to saving the entity the persist method will also assign a value to the field of the entity annotated with the @Id annotation (according to the configured @GenerationType). This behaviour is convenient but can be a bit tricky to test. To illustrate this lets use the following example.

Defining the problem

Lets say that we have a simple entity that represents a payment:

@Entity
public class Payment {
  @Id
  private Integer id;
  private BigDecimal amount;
  // getters and setters omitted
}

and that we want to test the method that inserts a new payment in to the database:

pubic class PaymentRepository {

  @PersistentContext
  EntityManager entityManager;

  @Inject
  EventRegistry eventRegistry;

  public void savePayment(Payment payment) {
    Preconditions.checkArgument(payment.getId() == null,
      "Must be a new payment");
    entityManager.persist(payment);
    eventRegistry.register(
      new SuccessfulPaymentEvent(payment.getId()));
  }
}

The above method checks if the payment is new, then saves the payment and then raises an appropriate event (SuccessfulPaymentEvent). For this method we would probably want to create two test cases. One when we try to create a payment that already exists and one for the case where the payment is successfully saved.

The first test case is trivial to implement:

@Test(expected=IllegalArgumentException.class)
public void failure_when_payment_exists() {
  PaymentRepository paymentRepo = new PaymentRepository();
  Payment payment = new Payment();
  payment.setId(10);

  paymentRepo.savePayment(payment);
}

In the second test case one of the things we want to verify is that the SuccessfulPaymentEvent is created with the correct payment id. To do so we have to assert that it is not null and also make sure it is the same as the id of the saved payment. This is not straightforward to implement because of the precondition check at the first line of the method. If we didn’t have the precondition then we could implement this using a Mockito ArgumentCaptor and pass a Payment with the id field already set to the savePayment(..) method. To illustrate the problem lets write a failing test and get an idea of what we are trying to achieve.

@Test
public void register_event_when_the_payment_is_successfully_saved() {
  PaymentRepository paymentRepo = new PaymentRepository();
  paymentRepo.entityManager = Mockito.mock(EntityManager.class);
  paymentRepo.eventRegistry = Mockito.mock(EventRegistry.class);
  ArgumentCaptor<SuccessfulPaymentEvent> argument =
    ArgumentCaptor.forClass(SuccessfulPaymentEvent.class);

  Payment payment = new Payment();
  paymentRepo.savePayment(payment);

  Mockito.verify(paymentRepo.entityManager).persist(payment);
  Mockito.verify(paymentRepo.eventRegistry).register(argument.capture());

  Integer paymentId = argument.getValue().getPaymentId();
  Assert.assertNotNull(paymentId); // Fails here!
  Assert.assertEquals(payment.getId(), paymentId);
}

The above test will fail because the paymentId is null in the first assertion of the test. To overcome this obstacle what we have to do is find a way to inject a known id value to the payment object when we pass it to the persist method of EntityManager. We can do this by using another Mockito feature called Answer.

A solution using Mockito

An Answer can be used to specify the action that is executed and the return value that is returned when you interact with a Mockito mock. So for our problem all we have to do is create an Answer that extracts the payment object from arguments of the mock invocation and sets its id to a known value. Then we can use this known value to assert the id of the captured SuccessfulPaymentEvent.

import org.mockito.stubbing.Answer;
import org.mockito.invocation.InvocationOnMock;

public class PaymentIdSetter extends Answer<Void> {

  private Integer paymentId;

  public PaymentIdSetter(Integer paymentId) {
    this.paymentId = paymentId;
  }

  @Override
  public Void answer(InvocationOnMock invocation) throws Throwable {
    Assert.assertEquals(1, invocation.getArguments().length);
    Payment payment = (Payment) invocation.getArguments()[0];
    payment.setId(paymentId);
    return null;
  }
}

Below we change our test case to use the PaymentIdSetter answer.

@Test
public void register_event_when_the_payment_is_successfully_saved() {
  Integer expectedId = 15;
  PaymentRepository paymentRepo = new PaymentRepository();

  paymentRepo.entityManager = Mockito.mock(EntityManager.class);
  Mockito.doAnswer(new PaymentIdSetter(expectedId))
    .when(paymentRepo.entityManager.persist(Matchers.any(Payment.class)));

  paymentRepo.eventRegistry = Mockito.mock(EventRegistry.class);
  ArgumentCaptor<SuccessfulPaymentEvent> argument =
    ArgumentCaptor.forClass(SuccessfulPaymentEvent.class);

  Payment payment = new Payment();
  paymentRepo.savePayment(payment);

  Mockito.verify(paymentRepo.entityManager).persist(payment);
  Mockito.verify(paymentRepo.eventRegistry).register(argument.capture());

  Integer paymentId = argument.getValue().getPaymentId();
  Assert.assertNotNull(paymentId);
  Assert.assertEquals(expectedId, paymentId);
}

Generalizing PaymentIdSetter

The test case required a bit more setup that is usual but it is able to verify the desired expectation. To improve this further we can easily change the PaymentIdSetter to support all the entity classes of our project either by using reflection to set the id value of the entities or by using a common parent interface (or abstract class) with a setId(..) method for our entity classes.

A possible implementation of the common parent interface option is the following:

public interface IdentifiableEntity {
  Integer getId();
  void setId(Integer id);
}

@Entity
public class Payment implements IdentifiableEntity {
  @Id
  private Integer id;
  private BigDecimal amount;

  public void setId(Integer id) {
    this.id = id;
  }

  // other getters and setters omitted
}

public class EntityIdSetter extends Answer<Void> {

  private Integer id;

  public EntityIdSetter(Integer id) {
    this.id = id;
  }

  @Override
  public Void answer(InvocationOnMock invocation) throws Throwable {
    Assert.assertEquals(1, invocation.getArguments().length);
    IdentifiableEntity entity = (IdentifiableEntity) invocation.getArguments()[0];
    entity.setId(id);
    return null;
  }
}

Conclusion

In this post we have seen how using a custom Mockito Answer enables us to effectively test code that invokes the EntityManager.persist(..) method. Of course this mechanism is not only applicable to JPA’s persist method and applies to testing code that calls methods that modify their parameters. JPA just provided us a nice example to explore this technique.