lxq/mydreamjob-server/src/test/java/com/mydreamjob/auth/company/CompanyRegistrationConcurrencyMySqlTest.java

package com.mydreamjob.auth.company;

import com.mydreamjob.audit.AuditContext;
import com.mydreamjob.auth.*;
import com.mydreamjob.company.CompanyRepository;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.boot.test.mock.mockito.MockBean;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.test.context.ActiveProfiles;
import org.springframework.test.context.DynamicPropertyRegistry;
import org.springframework.test.context.DynamicPropertySource;
import org.testcontainers.containers.MySQLContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;

import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.when;

@Testcontainers(disabledWithoutDocker = false)
@SpringBootTest
@ActiveProfiles("mysql-it")
class CompanyRegistrationConcurrencyMySqlTest {
    private static final String PHONE = "19900000321";
    private static final String CODE = "123456";

    @Container
    static final MySQLContainer<?> MYSQL = new MySQLContainer<>("mysql:8.4")
            .withDatabaseName("company_registration_concurrency")
            .withUsername("company_registration_test")
            .withPassword("company-registration-test");

    @DynamicPropertySource
    static void properties(DynamicPropertyRegistry registry) {
        registry.add("spring.datasource.url", MYSQL::getJdbcUrl);
        registry.add("spring.datasource.username", MYSQL::getUsername);
        registry.add("spring.datasource.password", MYSQL::getPassword);
        registry.add("app.company-registration.enabled", () -> "true");
    }

    @Autowired CompanyRegistrationService service;
    @Autowired SmsVerificationCodeRepository codes;
    @Autowired UserRepository users;
    @Autowired CompanyRepository companies;
    @Autowired JdbcTemplate jdbc;
    @MockBean RedisTemplate<String, Object> redis;

    @AfterEach
    void cleanup() {
        jdbc.update("DELETE FROM auth_tokens WHERE user_id IN (SELECT id FROM users WHERE phone = ?)", PHONE);
        jdbc.update("DELETE FROM companies WHERE owner_user_id IN (SELECT id FROM users WHERE phone = ?)", PHONE);
        jdbc.update("DELETE FROM user_roles WHERE user_id IN (SELECT id FROM users WHERE phone = ?)", PHONE);
        jdbc.update("DELETE FROM users WHERE phone = ?", PHONE);
        jdbc.update("DELETE FROM sms_verification_codes WHERE phone = ?", PHONE);
    }

    @Test
    void oneCodeCanCreateOnlyOneCompanyIdentityWithoutTimingSleeps() throws Exception {
        when(redis.hasKey(anyString())).thenReturn(false);
        codes.saveAndFlush(SmsVerificationCode.builder().phone(PHONE).purpose("company-register")
                .code(CODE).attempts(0).maxAttempts(5).used(false)
                .createdAt(LocalDateTime.now()).expiredAt(LocalDateTime.now().plusMinutes(5)).build());

        ExecutorService pool = Executors.newFixedThreadPool(2);
        CountDownLatch ready = new CountDownLatch(2);
        CountDownLatch start = new CountDownLatch(1);
        List<Future<Object>> futures = new ArrayList<>();
        for (int i = 0; i < 2; i++) {
            futures.add(pool.submit(() -> {
                ready.countDown();
                assertThat(start.await(10, TimeUnit.SECONDS)).isTrue();
                try {
                    return service.register(request(), AuditContext.builder().build());
                } catch (Throwable error) {
                    return error;
                }
            }));
        }
        assertThat(ready.await(10, TimeUnit.SECONDS)).isTrue();
        start.countDown();
        Object first = futures.get(0).get(30, TimeUnit.SECONDS);
        Object second = futures.get(1).get(30, TimeUnit.SECONDS);
        pool.shutdownNow();

        long successes = List.of(first, second).stream().filter(CompanyRegistrationResponse.class::isInstance).count();
        assertThat(successes).isEqualTo(1);
        assertThat(users.findAll().stream().filter(u -> PHONE.equals(u.getPhone())).count()).isEqualTo(1);
        Long userId = users.findByPhone(PHONE).orElseThrow().getId();
        assertThat(jdbc.queryForList("SELECT role FROM user_roles WHERE user_id = ?", String.class, userId))
                .containsExactly("COMPANY");
        assertThat(companies.findByOwnerUserId(userId)).isPresent();
        assertThat(jdbc.queryForObject("SELECT COUNT(*) FROM auth_tokens WHERE user_id = ?", Long.class, userId))
                .isEqualTo(1L);
        assertThat(List.of(first, second).stream().filter(Throwable.class::isInstance)
                .map(Object::toString).allMatch(text -> !text.contains("SQL") && !text.contains("phone"))).isTrue();
    }

    private CompanyRegistrationRequest request() {
        CompanyRegistrationRequest request = new CompanyRegistrationRequest();
        request.setPhone(PHONE); request.setCode(CODE); request.setPassword("Company-strong-password-2026");
        CompanyRegistrationRequest.CompanyDetails details = new CompanyRegistrationRequest.CompanyDetails();
        details.setName("并发企业"); details.setIndustry("互联网"); details.setDescription("简介"); details.setAddress("地址");
        request.setCompany(details);
        return request;
    }
}