Seeding Cosmos DB Data to run Spring Boot 2 Integration Tests

1. OVERVIEW

Let’s say you are deploying your Spring Boot RESTful applications to Azure.

Some of these Spring Boot applications might have been modernization rewrites to use Cosmos DB instead of relational databases.

You might have even added support to write dynamic Cosmos DB queries using Spring Data Cosmos. And let’s also assume you wrote unit tests for REST controllers, business logic, and utility classes.

Now you need to write integration tests to verify the interaction between different parts of the system work well, including retrieving from, and storing to, a NoSQL database like Cosmos DB.

This blog post shows you how to write a custom Spring’s TestExecutionListener to seed data in a Cosmos database container. Each Spring Boot integration test will run starting from a known Cosmos container state, so that you won’t need to force the tests to run in a specific order.

Seed data for integration tests in Spring Boot and Azure Cosmos DB applications

Tip: You might want to take a look at the recent Code and Configuration Snippets / Recipes / Tips section, which helps you with common tasks without reading through a full blog post.

2. COSMOS DB EMULATOR

This blog post uses Microsoft’s azure-cosmos-emulator Docker image.

It assumes your Spring Boot application already connects to a Cosmos DB container, located either in an emulator Docker container, in the native Windows emulator, or hosted in Azure.

Tip:
Read on if you are interested in learning how to extract a certificate and add it to a Truststore programmatically for your Spring Boot applications to connect to the Cosmos DB emulator Docker container.

You can also go over these two relevant code snippets:
Create KeyStore / Truststore Programmatically.
Add Cosmos DB emulator Certificate to TrustStore via Bash Shell Script.

3. MAVEN DEPENDENCIES

pom.xml:

...
<parent>
  <groupId>org.springframework.boot</groupId>
  <artifactId>spring-boot-starter-parent</artifactId>
  <version>2.7.14</version>
  <relativePath />
</parent>

<properties>
  <java.version>11</java.version>
  <spring-cloud-azure.version>4.10.0</spring-cloud-azure.version>
...
</properties>

<dependencies>
  <dependency>
    <groupId>com.azure.spring</groupId>
    <artifactId>spring-cloud-azure-starter-data-cosmos</artifactId>
  </dependency>
...
  <dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-test</artifactId>
    <scope>test</scope>
  </dependency>
...
<dependencyManagement>
  <dependencies>
    <dependency>
      <groupId>com.azure.spring</groupId>
      <artifactId>spring-cloud-azure-dependencies</artifactId>
      <version>${spring-cloud-azure.version}</version>
      <type>pom</type>
      <scope>import</scope>
    </dependency>
  </dependencies>
</dependencyManagement>

This tutorial uses Spring Boot version 2.7.14.
Azure BOM’s spring-cloud-azure-dependencies version is set to 4.10.0. The highest supported version at the time of this writing that you could use with Spring Boot 2.7.x. It manages a number of Spring, Azure-related libraries.

Important: Check this table to make sure you are not mismatching Spring Boot and Spring Cloud Azure Dependencies versions.

spring-cloud-azure-starter-data-cosmos dependency is managed by spring-cloud-azure-dependencies.
spring-boot-starter-test dependency is managed by a Spring Boot dependency transitively.

4. CONTAINER DOCUMENTS

User.java:

@Container(containerName = "users")
public class User {

  @Id
  private String id;
  private String firstName;
  private String lastName;
  private String address;
// ...
}

The name of the container space in the Cosmos database is users. It’ll store JSON documents corresponding to User.java objects.

5. TEST CONFIGURATION

src/test/resources/application-integration-test.yml:

spring:
  cloud:
    azure:
      cosmos:
        endpoint: https://localhost:8081
        key: C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==
        database: sampleIntegrationTestsDB

Make sure your integration tests use a Cosmos database different than what you use for development.
This blog post uses Microsoft’s azure-cosmos-emulator Docker image.
You could use a Cosmos database just for integration tests hosted in Azure Cloud.
You could also use Testcontainers.

Tip: Read on to learn:

How to use Testcontainers to write Integration Tests for Cosmos DB applications.

Reuse Testcontainers initialization and configuration code with JUnit 5 Extension Callbacks in your Spring Boot Integration Tests.

You would get the Cosmos key from the emulator home page.

Cosmos DB Emulator Key and Connection Strings

6. TESTEXECUTIONLISTENER SUPPORTING CODE

This custom Spring’s TestExecutionListener implementation hooks into the integration test life cycle to seed Cosmos data before running each test, and to remove it after each test completes.

CosmosDataTestExecutionListener.java:

public class CosmosDataTestExecutionListener extends AbstractTestExecutionListener {

  @Override
  public void beforeTestMethod(TestContext testContext) throws Exception {
    super.beforeTestMethod(testContext);
    this.executeCosmosOperation(testContext, ExecutionPhase.BEFORE_TEST_METHOD);
  }

  @Override
  public void afterTestMethod(TestContext testContext) throws Exception {
    RuntimeException resultingException = null;
    try {
      this.executeCosmosOperation(testContext, ExecutionPhase.AFTER_TEST_METHOD);
    } catch (RuntimeException ex) {
      resultingException = ex;
      log.warn("Swallowed exception to continue with the afterTestMethod() execution chain", ex);
    }
    super.afterTestMethod(testContext);
    if (resultingException != null) {
      throw resultingException;
    }
  }

  private void executeCosmosOperation(TestContext testContext, ExecutionPhase executionPhase) {
    boolean classLevel = false;

    Set<Cosmos> cosmosAnnotations = AnnotatedElementUtils.getMergedRepeatableAnnotations(
      testContext.getTestMethod(), Cosmos.class, CosmosGroup.class
    );
    if (cosmosAnnotations.isEmpty()) {
      cosmosAnnotations = AnnotatedElementUtils.getMergedRepeatableAnnotations(
        testContext.getTestClass(), Cosmos.class, CosmosGroup.class
      );
      if (!cosmosAnnotations.isEmpty()) {
        classLevel = true;
      }
    }

    this.executeCosmosOperation(cosmosAnnotations, testContext, executionPhase, classLevel);
  }

  private void executeCosmosOperation(Set<Cosmos> cosmosAnnotations, TestContext testContext,
      ExecutionPhase executionPhase, boolean classLevel) {

    cosmosAnnotations.forEach(cosmosAnnotation ->
      this.executeCosmosOperation(cosmosAnnotation, testContext, executionPhase, classLevel)
    );
  }

  private void executeCosmosOperation(Cosmos cosmosAnnotation, TestContext testContext,
      ExecutionPhase executionPhase, boolean classLevel) {
// ...

    ResourceCosmosDBPopulator populator = new ResourceCosmosDBPopulator(dataResource, containerClass, containerName,
      cosmosTemplate
    );
    if (cosmosAnnotation.executionPhase() == ExecutionPhase.BEFORE_TEST_METHOD) {
      populator.populate();
    } else {
      populator.deleteAll();
    }
  }

// ...
}

beforeTestMethod() method first follows the TestExecutionListener chain by calling super.beforeTestMethod(…).
It then executes a common executeCosmosOperation(…) passing the BEFORE_TEST_METHOD execution phase.

Multiple executeCosmosOperation() overloaded methods:
1. Retrieves and merges @Cosmos configuration annotations.
  It first tries to use the @CosmosGroup configuration from the current test method (line 38), and if it’s not found, it then tries the one at class level (line 9).
2. Loops through the merged @Cosmos annotations and calls another overloaded method passing the test execution phase and each @Cosmos annotation.
3. Retrieves data files, Cosmos DB container class and name from the @Cosmos configuration annotation, the ReactiveCosmosTemplate cosmosTemplate bean, and delegates inserting or removing data depending on the current test execution phase to ResourceCosmosDBPopulator.java.
afterTestMethod() is similar to beforeTestMethod().
Just that it first executes the same executeCosmosOperation(…) passing the AFTER_TEST_METHOD execution phase, but catching any exception so that it follows the super.afterTestMethod(…)’s TestExecutionListener chain to tear down or release other resources.

7. COSMOS DB DATA POPULATOR

CosmosDataTestExecutionListener delegates populating and deleting data to this class.

ResourceCosmosDBPopulator.java:

public class ResourceCosmosDBPopulator implements CosmosDBPopulator {

  private ObjectMapper objectMapper = new ObjectMapper();

  private final Resource data;
  private final Class<?> containerClass;
  private final String containerName;
  private final ReactiveCosmosTemplate cosmosTemplate;
// ...

  @Override
  public void populate() throws DataAccessException {
    log.info("Using Cosmos data: " + ObjectUtils.nullSafeToString(this.data));

    JavaType type = this.objectMapper.getTypeFactory().constructCollectionType(List.class, this.containerClass);
    try {
      List<?> values = this.objectMapper.readValue(this.data.getFile(), type);
      values.forEach(value -> this.cosmosTemplate.insert(value).block());
    } catch (Exception ex) {
      throw new CannotReadDataException(this.data, ex);
    }
  }

  @Override
  public void deleteAll() throws DataAccessException {
    this.cosmosTemplate.deleteAll(this.containerName, this.containerClass).block();
  }
}

populate() method reads a resource data file, converts a JSON array to a POJO list, and uses the ReactiveCosmosTemplate method argument to insert each list element in a Cosmos DB container.
deleteAll() method empties the Cosmos DB container.

8. TEST DATA FILES

This is a sample JSON document to seed data in a Cosmos DB container.

You reference this file in a custom @Cosmos annotation discussed in the next section.

src/test/resources/data-1.json:

[
  {
    "id": "100",
    "firstName": "First-100",
    "lastName": "Last-100",
    "address": "100 Main St"
  },
  {
    "id": "101",
    "firstName": "First-101",
    "lastName": "Last-101",
    "address": "101 Main St"
  },
  {
    "id": "102",
    "firstName": "First-102",
    "lastName": "Last-102",
    "address": "102 Main St"
  }
]

9. INTEGRATION TEST CLASS

Let’s write an integration test class that uses everything laid out in this post.

UserControllerIntegrationTest.java:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
@ExtendWith(SpringExtension.class)
@SpringBootTest(classes = Application.class, webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
@ContextConfiguration(classes = { Application.class })
@ActiveProfiles("integration-test")
@TestExecutionListeners({
  DependencyInjectionTestExecutionListener.class,
  CosmosDataTestExecutionListener.class
})
@CosmosGroup({
  @Cosmos(data = "classpath:data-1.json", containerClass = User.class, executionPhase = ExecutionPhase.BEFORE_TEST_METHOD),
  @Cosmos(containerClass = User.class, containerName = "users", executionPhase = ExecutionPhase.AFTER_TEST_METHOD)
})
@Slf4j
public class UserControllerIntegrationTest {

  @LocalServerPort
  private int port;
// ...

  @Test
  public void shouldRetrieveUserWithIdEqualsTo100() {
    JsonPath jsonPath = RestAssured
      .given()
        .accept(ContentType.JSON)
      .when()
        .get("/api/users/100")
      .then()
        .statusCode(HttpStatus.OK.value())
        .contentType(ContentType.JSON)
        .extract()
        .jsonPath();

    Map<String, Object> actualUser = jsonPath.get("$");
    MatcherAssert.assertThat(actualUser.get("id"), Matchers.equalTo("100"));
    MatcherAssert.assertThat(actualUser.get("firstName"), Matchers.equalTo("First-100"));
  }

  @CosmosGroup({
    @Cosmos(data = "classpath:data-2.json", containerClass = User.class, executionPhase = ExecutionPhase.BEFORE_TEST_METHOD),
    @Cosmos(containerClass = User.class, containerName = "users", executionPhase = ExecutionPhase.AFTER_TEST_METHOD)
  })
  @Test
  public void shouldRetrieveAllUsers() {
    JsonPath jsonPath = RestAssured
      .given()
        .accept(ContentType.JSON)
      .when()
        .get("/api/users")
      .then()
        .statusCode(HttpStatus.OK.value())
        .contentType(ContentType.JSON)
        .extract()
        .jsonPath();

    List<Object> actualUsers = jsonPath.get("$");
    MatcherAssert.assertThat(actualUsers.size(), Matchers.equalTo(3));
  }
// ...
}

This integration test class uses JUnit 5 (brought in when including spring-boot-starter-test:2.7.14.

@ActiveProfiles annotation includes the integration-test profile for this integration test class to use application-integration-test.yml test properties file.

@TestExecutionListeners annotation includes:

DependencyInjectionTestExecutionListener because it updates:

@LocalServerPort
private int port;

for RestAssured to send requests to.

and the custom CosmosDataTestExecutionListener Spring’s TestExecutionListener.

CosmosDataTestExecutionListener is used in conjunction with @CosmosGroup and @Cosmos annotations.

@CosmosGroup and @Cosmos are custom configuration annotations. They are meant to work similarly to when you use @SqlGroup and @Sql to write integration tests using a relational database.

@CosmosGroup configures CosmosDataTestExecutionListener to seed data specified in data-1.json into the Cosmos container referenced in User’s @Container annotation before each test runs.

@CosmosGroup also configures CosmosDataTestExecutionListener to empty the container users after each test completes.

Notice how shouldRetrieveAllUsers() method also includes @CosmosGroup and @Cosmos configuration annotations.
The CosmosDataTestExecutionListener implementation allows to use the @CosmosGroup and @Cosmos annotations at the class level as well as at the individual test level.

Seed data for integration tests in Spring Boot and Azure Cosmos DB applications

10. CONCLUSION

This blog post covered how to write integration tests for Spring Boot applications that connect to Cosmos databases the same way you would when connecting them to relational databases.

You can write a custom TestExecutionListener that hooks into the integration test life cycle to seed a Cosmos DB container before running each test, and empties it after each test completes.
This approach ensures that each integration test runs from a known Cosmos database state without imposing tests execution order.

You could provision a Cosmos database to run integration tests the same way you would when using @SqlGroup, @Sql annotations to provision a relation database.

Thanks for reading and as always, feedback is very much appreciated. If you found this post helpful and would like to receive updates when content like this gets published, sign up to the newsletter.