package tech.easyflow.datacenter.connector.support;

import org.junit.Test;
import org.mockito.InOrder;
import org.mockito.Mockito;
import tech.easyflow.datacenter.connector.dialect.PostgresqlSqlDialect;
import tech.easyflow.datacenter.meta.entity.DatacenterSource;
import tech.easyflow.datacenter.meta.enums.DatacenterCapability;
import tech.easyflow.datacenter.meta.enums.DatacenterSourceType;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.util.EnumSet;

/**
 * PostgreSQL 服务端游标连接状态回归测试。
 */
public class PostgresqlStreamingConnectorTest {

    /**
     * 验证自动提交连接进入游标事务并在消费完成后恢复。
     *
     * @throws Exception JDBC 模拟调用失败时抛出
     */
    @Test
    public void shouldRestoreAutoCommitAfterStreaming()
            throws Exception {
        Connection connection =
                Mockito.mock(Connection.class);
        PreparedStatement statement =
                Mockito.mock(
                        PreparedStatement.class);
        ResultSet resultSet =
                Mockito.mock(ResultSet.class);
        ResultSetMetaData metaData =
                Mockito.mock(
                        ResultSetMetaData.class);
        Mockito.when(connection.getAutoCommit())
                .thenReturn(true);
        Mockito.when(connection.prepareStatement(
                        "SELECT id FROM sample",
                        ResultSet.TYPE_FORWARD_ONLY,
                        ResultSet.CONCUR_READ_ONLY))
                .thenReturn(statement);
        Mockito.when(statement.executeQuery())
                .thenReturn(resultSet);
        Mockito.when(resultSet.getMetaData())
                .thenReturn(metaData);
        Mockito.when(resultSet.next())
                .thenReturn(false);
        TestConnector connector =
                new TestConnector(connection);

        connector.consumeBySql(
                new DatacenterSource(),
                "SELECT id FROM sample",
                512,
                row -> {
                });

        InOrder order = Mockito.inOrder(
                connection,
                statement,
                resultSet);
        order.verify(connection)
                .getAutoCommit();
        order.verify(connection)
                .setAutoCommit(false);
        order.verify(connection)
                .prepareStatement(
                        "SELECT id FROM sample",
                        ResultSet.TYPE_FORWARD_ONLY,
                        ResultSet.CONCUR_READ_ONLY);
        order.verify(statement)
                .setFetchSize(512);
        order.verify(statement)
                .executeQuery();
        order.verify(resultSet)
                .close();
        order.verify(statement)
                .close();
        order.verify(connection)
                .rollback();
        order.verify(connection)
                .setAutoCommit(true);
    }

    /**
     * 使用测试连接执行 PostgreSQL 查询。
     */
    private static final class TestConnector
            extends AbstractJdbcConnector {

        private final Connection connection;

        /**
         * 创建测试连接器。
         *
         * @param connection 测试 JDBC 连接
         */
        private TestConnector(
                Connection connection) {
            super(
                    DatacenterSourceType.POSTGRESQL,
                    new PostgresqlSqlDialect(),
                    EnumSet.of(
                            DatacenterCapability.READ_QUERY));
            this.connection = connection;
        }

        /**
         * {@inheritDoc}
         */
        @Override
        protected <T> T withConnection(
                DatacenterSource source,
                boolean cacheable,
                JdbcCallback<T> callback)
                throws Exception {
            return callback.apply(connection);
        }
    }
}
