package com.easyagents.federation.sql.execute;

import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import com.easyagents.federation.sql.source.SourceId;
import java.time.Duration;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import org.junit.Assert;
import org.junit.Test;

/**
 * 本地查询准入控制器关闭与等待唤醒测试。
 */
public class LocalFederationQueryAdmissionControllerTest {

    /**
     * 验证关闭控制器会唤醒等待线程，并拒绝后续准入。
     *
     * @throws Exception 并发测试失败
     */
    @Test
    public void shouldWakeAndRejectWaitersAfterClose() throws Exception {
        LocalFederationQueryAdmissionController controller =
            new LocalFederationQueryAdmissionController(1);
        FederationQueryPermit first = controller.acquire(
            new SourceId("main"),
            new QueryId("first"),
            Duration.ZERO
        );
        ExecutorService executor = Executors.newSingleThreadExecutor();
        try {
            Future<FederationSqlErrorCode> waiting = executor.submit(() -> {
                try {
                    controller.acquire(
                        new SourceId("main"),
                        new QueryId("waiting"),
                        Duration.ofSeconds(5)
                    );
                    return null;
                } catch (FederationSqlException exception) {
                    return exception.errorCode();
                }
            });
            controller.close();
            Assert.assertEquals(FederationSqlErrorCode.ENGINE_CLOSED, waiting.get(2, TimeUnit.SECONDS));
            try {
                controller.acquire(new SourceId("main"), new QueryId("late"), Duration.ZERO);
                Assert.fail("closed controller should reject admission");
            } catch (FederationSqlException exception) {
                Assert.assertEquals(FederationSqlErrorCode.ENGINE_CLOSED, exception.errorCode());
            }
        } finally {
            first.close();
            controller.close();
            executor.shutdownNow();
        }
    }
}
