Testing a Spring Boot Backend with Multi tenancy fails because Session has old tenant

I have a Spring Boot Application using Hibernate in a multi tenancy setup with shared database and shared schema setup. The schema is created using liquibase.

My problem is, that the application works fine when I start it normally. In tests however the application fails because the database session still has the default tenant.

I am testing against an H2 in memory database.

The relevant classes are these:

// holds the tenant id
public class TenantContext {

    private TenantContext() {}

    private static final ThreadLocal<Integer> CURRENT_TENANT = new ThreadLocal<>();

    public static Integer getCurrentTenant() {
        return CURRENT_TENANT.get();
    }

    public static void setCurrentTenant(Integer tenant) {
        CURRENT_TENANT.set(tenant);
    }

    public static void reset() {
        CURRENT_TENANT.remove();
    }
}

// provides the tenant id to hibernate
@Component
public class TenantResolver implements CurrentTenantIdentifierResolver<Integer>, HibernatePropertiesCustomizer {

    @Override
    public Integer resolveCurrentTenantIdentifier() {
        Integer tenant = TenantContext.getCurrentTenant();
        if(tenant == null) return -1;
        return tenant;
    }

    @Override
    public boolean validateExistingCurrentSessions() {
        return false;
    }

    @Override
    public void customize(Map<String, Object> hibernateProperties) {
        hibernateProperties.put(MultiTenancySettings.MULTI_TENANT_IDENTIFIER_RESOLVER, this);
    }
}

// Entity class
@Entity
@Getter
@Setter
@Builder(toBuilder = true)
@AllArgsConstructor
@NoArgsConstructor
@Table(schema = "MySchema", name = "MyTable")
public class MyEntity{

    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    @Column(name = "id")
    private Long id;

    @TenantId
    @Column(name = "tenant")
    private Integer tenant;
}

public interface MyEntityRepository extends CrudRepository<MyEntity, Long> {}

I don't implement MultiTenantConnectionProvider as I have a shared database and schema.

The test class looks like this:

@SpringBootTest
@Transactional
@AutoConfigureMockMvc
@ActiveProfiles(profiles = "h2")
class MyEntityRepositoryTest{
    @Autowired
    MyEntityRepository myEntityRepository;

    @Test
    void createEntityWithTenant() {
        TenantContext.setCurrentTenant(123);
        MyEntity myEntity = new MyEntity();
        MyEntity storedEntity = myEntityRepository.save(myEntity);
        assertEquals(123, storedEntity.getTenant());
    }
}

This fails, because myEntity.getTenant() returns -1, the default tenant. This is because the SessionImpl that executes the transaction still has the old tenant id from initialization.

// this is from debugging inside InsertCoordinator#preInsertInMemoryValueGeneration
// https://docs.jboss.org/hibernate/stable/core/javadocs/org/hibernate/persister/entity/mutation/InsertCoordinator.html#preInsertInMemoryValueGeneration(java.lang.Object%5B%5D,java.lang.Object,org.hibernate.engine.spi.SharedSessionContractImplementor)

session.getTenantIdentifierValue(); // returns -1

session.getSessionFactory()
       .getCurrentTenantIdentifierResolver()
       .resolveCurrentTenantIdentifier() // returns 123

session.getSessionFactory()
       .openSession()
       .getTenantIdentifierValue() // returns 123

I am clueless how to proceed. I think creating a new session that has the updated tenant id is possible, but how do I make myEntityRepository use this new session?

Removing the @Transactional annotation from class MyEntityRepositoryTest solved the issue.

@SpringBootTest
// @Transactional
@AutoConfigureMockMvc
@ActiveProfiles(profiles = "h2")
class MyEntityRepositoryTest{
    @Autowired
    MyEntityRepository myEntityRepository;

    @Test
    void createEntityWithTenant() {
        TenantContext.setCurrentTenant(123);
        MyEntity myEntity = new MyEntity();
        MyEntity storedEntity = myEntityRepository.save(myEntity);
        assertEquals(123, storedEntity.getTenant());
    }
}