/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.kafka.clients.admin.internals;

import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.admin.internals.AdminApiDriver.RequestSpec;
import org.apache.kafka.clients.admin.internals.AdminApiHandler.ApiResult;
import org.apache.kafka.clients.admin.internals.AdminApiLookupStrategy.LookupResult;
import org.apache.kafka.common.KafkaFuture;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.errors.DisconnectException;
import org.apache.kafka.common.errors.UnknownServerException;
import org.apache.kafka.common.errors.UnsupportedVersionException;
import org.apache.kafka.common.message.MetadataResponseData;
import org.apache.kafka.common.protocol.ApiKeys;
import org.apache.kafka.common.requests.AbstractRequest;
import org.apache.kafka.common.requests.AbstractResponse;
import org.apache.kafka.common.requests.FindCoordinatorRequest.NoBatchedFindCoordinatorsException;
import org.apache.kafka.common.requests.MetadataRequest;
import org.apache.kafka.common.requests.MetadataResponse;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.common.utils.MockTime;

import org.junit.jupiter.api.Test;

import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.OptionalInt;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;

import static java.util.Collections.emptyMap;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;

class AdminApiDriverTest {
    private static final int API_TIMEOUT_MS = 30000;
    private static final int RETRY_BACKOFF_MS = 100;
    private static final int RETRY_BACKOFF_MAX_MS = 1000;

    @Test
    public void testCoalescedLookup() {
        TestContext ctx = TestContext.dynamicMapped(map(
            "foo", "c1",
            "bar", "c1"
        ));

        Map<Set<String>, LookupResult<String>> lookupRequests = map(
            Set.of("foo", "bar"), mapped("foo", 1, "bar", 2)
        );

        ctx.poll(lookupRequests, emptyMap());

        Map<Set<String>, ApiResult<String, Long>> fulfillmentResults = map(
            Set.of("foo"), completed("foo", 15L),
            Set.of("bar"), completed("bar", 30L)
        );

        ctx.poll(emptyMap(), fulfillmentResults);

        ctx.poll(emptyMap(), emptyMap());
    }

    @Test
    public void testCoalescedFulfillment() {
        TestContext ctx = TestContext.dynamicMapped(map(
            "foo", "c1",
            "bar", "c2"
        ));

        Map<Set<String>, LookupResult<String>> lookupRequests = map(
            Set.of("foo"), mapped("foo", 1),
            Set.of("bar"), mapped("bar", 1)
        );

        ctx.poll(lookupRequests, emptyMap());

        Map<Set<String>, ApiResult<String, Long>> fulfillmentResults = map(
            Set.of("foo", "bar"), completed("foo", 15L, "bar", 30L)
        );

        ctx.poll(emptyMap(), fulfillmentResults);

        ctx.poll(emptyMap(), emptyMap());
    }

    @Test
    public void testKeyLookupFailure() {
        // Ensure that both generic failures and unhandled UnsupportedVersionExceptions (which could be specifically
        // handled in both the lookup and the fulfillment stages) result in the expected lookup failures.
        Exception[] keyLookupExceptions = new Exception[] {
            new UnknownServerException(), new UnsupportedVersionException("")
        };
        for (Exception keyLookupException : keyLookupExceptions) {
            TestContext ctx = TestContext.dynamicMapped(map(
                "foo", "c1",
                "bar", "c2"
            ));

            Map<Set<String>, LookupResult<String>> lookupRequests = map(
                Set.of("foo"), failedLookup("foo", keyLookupException),
                Set.of("bar"), mapped("bar", 1)
            );

            ctx.poll(lookupRequests, emptyMap());

            Map<Set<String>, ApiResult<String, Long>> fulfillmentResults = map(
                Set.of("bar"), completed("bar", 30L)
            );

            ctx.poll(emptyMap(), fulfillmentResults);

            ctx.poll(emptyMap(), emptyMap());
        }
    }

    @Test
    public void testKeyLookupRetry() {
        TestContext ctx = TestContext.dynamicMapped(map(
            "foo", "c1",
            "bar", "c2"
        ));

        Map<Set<String>, LookupResult<String>> lookupRequests = map(
            Set.of("foo"), emptyLookup(),
            Set.of("bar"), mapped("bar", 1)
        );

        ctx.poll(lookupRequests, emptyMap());

        Map<Set<String>, LookupResult<String>> fooRetry = map(
            Set.of("foo"), mapped("foo", 1)
        );

        Map<Set<String>, ApiResult<String, Long>> barFulfillment = map(
            Set.of("bar"), completed("bar", 30L)
        );

        ctx.poll(fooRetry, barFulfillment);

        Map<Set<String>, ApiResult<String, Long>> fooFulfillment = map(
            Set.of("foo"), completed("foo", 15L)
        );

        ctx.poll(emptyMap(), fooFulfillment);

        ctx.poll(emptyMap(), emptyMap());
    }

    @Test
    public void testStaticMapping() {
        TestContext ctx = TestContext.staticMapped(map(
            "foo", 0,
            "bar", 1,
            "baz", 1
        ));

        Map<Set<String>, ApiResult<String, Long>> fulfillmentResults = map(
            Set.of("foo"), completed("foo", 15L),
            Set.of("bar", "baz"), completed("bar", 30L, "baz", 45L)
        );

        ctx.poll(emptyMap(), fulfillmentResults);

        ctx.poll(emptyMap(), emptyMap());
    }

    @Test
    public void testFulfillmentFailure() {
        TestContext ctx = TestContext.staticMapped(map(
            "foo", 0,
            "bar", 1,
            "baz", 1
        ));

        Map<Set<String>, ApiResult<String, Long>> fulfillmentResults = map(
            Set.of("foo"), failed("foo", new UnknownServerException()),
            Set.of("bar", "baz"), completed("bar", 30L, "baz", 45L)
        );

        ctx.poll(emptyMap(), fulfillmentResults);

        ctx.poll(emptyMap(), emptyMap());
    }

    @Test
    public void testFulfillmentRetry() {
        TestContext ctx = TestContext.staticMapped(map(
            "foo", 0,
            "bar", 1,
            "baz", 1
        ));

        Map<Set<String>, ApiResult<String, Long>> fulfillmentResults = map(
            Set.of("foo"), completed("foo", 15L),
            Set.of("bar", "baz"), completed("bar", 30L)
        );

        ctx.poll(emptyMap(), fulfillmentResults);

        Map<Set<String>, ApiResult<String, Long>> bazRetry = map(
            Set.of("baz"), completed("baz", 45L)
        );

        ctx.poll(emptyMap(), bazRetry);

        ctx.poll(emptyMap(), emptyMap());
    }

    @Test
    public void testFulfillmentUnmapping() {
        TestContext ctx = TestContext.dynamicMapped(map(
            "foo", "c1",
            "bar", "c2"
        ));

        Map<Set<String>, LookupResult<String>> lookupRequests = map(
            Set.of("foo"), mapped("foo", 0),
            Set.of("bar"), mapped("bar", 1)
        );

        ctx.poll(lookupRequests, emptyMap());

        Map<Set<String>, ApiResult<String, Long>> fulfillmentResults = map(
            Set.of("foo"), completed("foo", 15L),
            Set.of("bar"), unmapped("bar")
        );

        ctx.poll(emptyMap(), fulfillmentResults);

        Map<Set<String>, LookupResult<String>> barLookupRetry = map(
            Set.of("bar"), mapped("bar", 1)
        );

        ctx.poll(barLookupRetry, emptyMap());

        Map<Set<String>, ApiResult<String, Long>> barFulfillRetry = map(
            Set.of("bar"), completed("bar", 30L)
        );

        ctx.poll(emptyMap(), barFulfillRetry);

        ctx.poll(emptyMap(), emptyMap());
    }

    @Test
    public void testFulfillmentFailureUnsupportedVersion() {
        TestContext ctx = TestContext.staticMapped(map(
            "foo", 0,
            "bar", 1,
            "baz", 1
        ));

        Map<Set<String>, ApiResult<String, Long>> fulfillmentResults = map(
            Set.of("foo"), failed("foo", new UnsupportedVersionException("")),
            Set.of("bar", "baz"), completed("bar", 30L, "baz", 45L)
        );

        ctx.poll(emptyMap(), fulfillmentResults);
        ctx.poll(emptyMap(), emptyMap());
    }

    @Test
    public void testFulfillmentRetriableUnsupportedVersion() {
        TestContext ctx = TestContext.staticMapped(map(
            "foo", 0,
            "bar", 1,
            "baz", 2
        ));

        ctx.handler.addRetriableUnsupportedVersionKey("foo");
        // The mapped ApiResults are only used in the onResponse/handleResponse path - anything that needs
        // to be handled in the onFailure path needs to be manually set up later.
        ctx.handler.expectRequest(Set.of("foo"), failed("foo", new UnsupportedVersionException("")));
        ctx.handler.expectRequest(Set.of("bar"), failed("bar", new UnsupportedVersionException("")));
        ctx.handler.expectRequest(Set.of("baz"), completed("baz", 45L));
        // Setting up specific fulfillment stage executions requires polling the driver in order to obtain
        // the request specs needed for the onResponse/onFailure callbacks.
        List<RequestSpec<String>> requestSpecs = ctx.driver.poll();

        requestSpecs.forEach(requestSpec -> {
            if (requestSpec.keys.contains("foo") || requestSpec.keys.contains("bar")) {
                ctx.driver.onFailure(ctx.time.milliseconds(), requestSpec, new UnsupportedVersionException(""));
            } else {
                ctx.driver.onResponse(
                    ctx.time.milliseconds(),
                    requestSpec,
                    new MetadataResponse(new MetadataResponseData(), ApiKeys.METADATA.latestVersion()),
                    Node.noNode());
            }
        });
        // Verify retry for "foo" but not for "bar" or "baz"
        ctx.poll(emptyMap(), map(
            Set.of("foo"), failed("foo", new UnsupportedVersionException(""))
        ));
        ctx.poll(emptyMap(), emptyMap());
    }

    @Test
    public void testRecoalescedLookup() {
        TestContext ctx = TestContext.dynamicMapped(map(
            "foo", "c1",
            "bar", "c1"
        ));

        Map<Set<String>, LookupResult<String>> lookupRequests = map(
            Set.of("foo", "bar"), mapped("foo", 1, "bar", 2)
        );

        ctx.poll(lookupRequests, emptyMap());

        Map<Set<String>, ApiResult<String, Long>> fulfillment = map(
            Set.of("foo"), unmapped("foo"),
            Set.of("bar"), unmapped("bar")
        );

        ctx.poll(emptyMap(), fulfillment);

        Map<Set<String>, LookupResult<String>> retryLookupRequests = map(
            Set.of("foo", "bar"), mapped("foo", 3, "bar", 3)
        );

        ctx.poll(retryLookupRequests, emptyMap());

        Map<Set<String>, ApiResult<String, Long>> retryFulfillment = map(
            Set.of("foo", "bar"), completed("foo", 15L, "bar", 30L)
        );

        ctx.poll(emptyMap(), retryFulfillment);

        ctx.poll(emptyMap(), emptyMap());
    }

    @Test
    public void testRetryLookupAfterDisconnect() {
        TestContext ctx = TestContext.dynamicMapped(map(
            "foo", "c1"
        ));

        int initialLeaderId = 1;

        Map<Set<String>, LookupResult<String>> initialLookup = map(
            Set.of("foo"), mapped("foo", initialLeaderId)
        );

        ctx.poll(initialLookup, emptyMap());
        assertMappedKey(ctx, "foo", initialLeaderId);

        ctx.handler.expectRequest(Set.of("foo"), completed("foo", 15L));

        List<RequestSpec<String>> requestSpecs = ctx.driver.poll();
        assertEquals(1, requestSpecs.size());

        RequestSpec<String> requestSpec = requestSpecs.get(0);
        assertEquals(OptionalInt.of(initialLeaderId), requestSpec.scope.destinationBrokerId());

        ctx.driver.onFailure(ctx.time.milliseconds(), requestSpec, new DisconnectException());
        assertUnmappedKey(ctx, "foo");

        int retryLeaderId = 2;

        ctx.lookupStrategy().expectLookup(Set.of("foo"), mapped("foo", retryLeaderId));
        List<RequestSpec<String>> retryLookupSpecs = ctx.driver.poll();
        assertEquals(1, retryLookupSpecs.size());

        RequestSpec<String> retryLookupSpec = retryLookupSpecs.get(0);
        assertEquals(ctx.time.milliseconds(), retryLookupSpec.nextAllowedTryMs);
        assertEquals(1, retryLookupSpec.tries);
    }

    @Test
    public void testRetryLookupAndDisableBatchAfterNoBatchedFindCoordinatorsException() {
        MockTime time = new MockTime();
        LogContext lc = new LogContext();
        Set<String> groupIds = Set.of("g1", "g2");
        DeleteConsumerGroupsHandler handler = new DeleteConsumerGroupsHandler(lc);
        AdminApiFuture<CoordinatorKey, Void> future = AdminApiFuture.forKeys(
                groupIds.stream().map(CoordinatorKey::byGroupId).collect(Collectors.toSet()));

        AdminApiDriver<CoordinatorKey, Void> driver = new AdminApiDriver<>(
            handler,
            future,
            time.milliseconds() + API_TIMEOUT_MS,
            RETRY_BACKOFF_MS,
            RETRY_BACKOFF_MAX_MS,
            new LogContext()
        );

        assertTrue(((CoordinatorStrategy) handler.lookupStrategy()).batch);
        List<RequestSpec<CoordinatorKey>> requestSpecs = driver.poll();
        // Expect CoordinatorStrategy to try resolving all coordinators in a single request
        assertEquals(1, requestSpecs.size());

        RequestSpec<CoordinatorKey> requestSpec = requestSpecs.get(0);
        driver.onFailure(time.milliseconds(), requestSpec, new NoBatchedFindCoordinatorsException("message"));
        assertFalse(((CoordinatorStrategy) handler.lookupStrategy()).batch);

        // Batching is now disabled, so we now have a request per groupId
        List<RequestSpec<CoordinatorKey>> retryLookupSpecs = driver.poll();
        assertEquals(groupIds.size(), retryLookupSpecs.size());
        // These new requests are treated a new requests and not retries
        for (RequestSpec<CoordinatorKey> retryLookupSpec : retryLookupSpecs) {
            assertEquals(0, retryLookupSpec.nextAllowedTryMs);
            assertEquals(0, retryLookupSpec.tries);
        }
    }

    @Test
    public void testCoalescedStaticAndDynamicFulfillment() {
        Map<String, String> dynamicMapping = map(
            "foo", "c1"
        );

        Map<String, Integer> staticMapping = map(
            "bar", 1
        );

        TestContext ctx = new TestContext(
            staticMapping,
            dynamicMapping
        );

        // Initially we expect a lookup for the dynamic key and a
        // fulfillment request for the static key
        LookupResult<String> lookupResult = mapped("foo", 1);
        ctx.lookupStrategy().expectLookup(
            Set.of("foo"), lookupResult
        );
        ctx.handler.expectRequest(
            Set.of("bar"), completed("bar", 10L)
        );

        List<RequestSpec<String>> requestSpecs = ctx.driver.poll();
        assertEquals(2, requestSpecs.size());

        RequestSpec<String> lookupSpec = requestSpecs.get(0);
        assertEquals(Set.of("foo"), lookupSpec.keys);
        ctx.assertLookupResponse(lookupSpec, lookupResult);

        // Receive a disconnect from the fulfillment request so that
        // we have an opportunity to coalesce the keys.
        RequestSpec<String> fulfillmentSpec = requestSpecs.get(1);
        assertEquals(Set.of("bar"), fulfillmentSpec.keys);
        ctx.driver.onFailure(ctx.time.milliseconds(), fulfillmentSpec, new DisconnectException());

        // Now we should get two fulfillment requests. One of them will
        // the coalesced dynamic and static keys for broker 1. The other
        // should contain the single dynamic key for broker 0.
        ctx.handler.reset();
        ctx.handler.expectRequest(
            Set.of("foo", "bar"), completed("foo", 15L, "bar", 30L)
        );

        List<RequestSpec<String>> coalescedSpecs = ctx.driver.poll();
        assertEquals(1, coalescedSpecs.size());
        RequestSpec<String> coalescedSpec = coalescedSpecs.get(0);
        assertEquals(Set.of("foo", "bar"), coalescedSpec.keys);

        // Disconnect in order to ensure that only the dynamic key is unmapped.
        // Then complete the remaining requests.
        ctx.driver.onFailure(ctx.time.milliseconds(), coalescedSpec, new DisconnectException());

        Map<Set<String>, LookupResult<String>> fooLookupRetry = map(
            Set.of("foo"), mapped("foo", 3)
        );
        Map<Set<String>, ApiResult<String, Long>> barFulfillmentRetry = map(
            Set.of("bar"), completed("bar", 30L)
        );
        ctx.poll(fooLookupRetry, barFulfillmentRetry);

        Map<Set<String>, ApiResult<String, Long>> fooFulfillmentRetry = map(
            Set.of("foo"), completed("foo", 15L)
        );
        ctx.poll(emptyMap(), fooFulfillmentRetry);
        ctx.poll(emptyMap(), emptyMap());
    }

    @Test
    public void testLookupRetryBookkeeping() {
        TestContext ctx = TestContext.dynamicMapped(map(
            "foo", "c1"
        ));

        LookupResult<String> emptyLookup = emptyLookup();
        ctx.lookupStrategy().expectLookup(Set.of("foo"), emptyLookup);

        List<RequestSpec<String>> requestSpecs = ctx.driver.poll();
        assertEquals(1, requestSpecs.size());

        RequestSpec<String> requestSpec = requestSpecs.get(0);
        assertEquals(0, requestSpec.tries);
        assertEquals(0L, requestSpec.nextAllowedTryMs);
        ctx.assertLookupResponse(requestSpec, emptyLookup);

        List<RequestSpec<String>> retrySpecs = ctx.driver.poll();
        assertEquals(1, retrySpecs.size());

        RequestSpec<String> retrySpec = retrySpecs.get(0);
        assertEquals(1, retrySpec.tries);
        assertEquals(ctx.time.milliseconds(), retrySpec.nextAllowedTryMs);
    }

    @Test
    public void testFulfillmentRetryBookkeeping() {
        TestContext ctx = TestContext.staticMapped(map("foo", 0));

        ApiResult<String, Long> emptyFulfillment = emptyFulfillment();
        ctx.handler.expectRequest(Set.of("foo"), emptyFulfillment);

        List<RequestSpec<String>> requestSpecs = ctx.driver.poll();
        assertEquals(1, requestSpecs.size());

        RequestSpec<String> requestSpec = requestSpecs.get(0);
        assertEquals(0, requestSpec.tries);
        assertEquals(0L, requestSpec.nextAllowedTryMs);
        ctx.assertResponse(requestSpec, emptyFulfillment, Node.noNode());

        List<RequestSpec<String>> retrySpecs = ctx.driver.poll();
        assertEquals(1, retrySpecs.size());

        RequestSpec<String> retrySpec = retrySpecs.get(0);
        assertEquals(1, retrySpec.tries);
        assertEquals(ctx.time.milliseconds(), retrySpec.nextAllowedTryMs,
                (long) (RETRY_BACKOFF_MS * CommonClientConfigs.RETRY_BACKOFF_JITTER));
    }

    private static void assertMappedKey(
        TestContext context,
        String key,
        Integer expectedBrokerId
    )  {
        OptionalInt brokerIdOpt = context.driver.keyToBrokerId(key);
        assertEquals(OptionalInt.of(expectedBrokerId), brokerIdOpt);
    }

    private static void assertUnmappedKey(
        TestContext context,
        String key
    ) {
        OptionalInt brokerIdOpt = context.driver.keyToBrokerId(key);
        assertEquals(OptionalInt.empty(), brokerIdOpt);
        KafkaFuture<Long> future = context.future.all().get(key);
        assertFalse(future.isDone());
    }

    private static void assertFailedKey(
        TestContext context,
        String key,
        Throwable expectedException
    ) {
        KafkaFuture<Long> future = context.future.all().get(key);
        assertTrue(future.isCompletedExceptionally());
        Throwable exception = assertThrows(ExecutionException.class, future::get);
        assertEquals(expectedException, exception.getCause());
    }

    private static void assertCompletedKey(
        TestContext context,
        String key,
        Long expected
    ) {
        KafkaFuture<Long> future = context.future.all().get(key);
        assertTrue(future.isDone());
        try {
            assertEquals(expected, future.get());
        } catch (Throwable t) {
            throw new RuntimeException(t);
        }
    }

    private static class MockRequestScope implements ApiRequestScope {
        private final OptionalInt destinationBrokerId;
        private final String id;

        private MockRequestScope(
            OptionalInt destinationBrokerId,
            String id
        ) {
            this.destinationBrokerId = destinationBrokerId;
            this.id = id;
        }

        @Override
        public OptionalInt destinationBrokerId() {
            return destinationBrokerId;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            MockRequestScope that = (MockRequestScope) o;
            return Objects.equals(destinationBrokerId, that.destinationBrokerId) &&
                Objects.equals(id, that.id);
        }

        @Override
        public int hashCode() {
            return Objects.hash(destinationBrokerId, id);
        }
    }

    private static class TestContext {
        private final MockTime time = new MockTime();
        private final MockAdminApiHandler<String, Long> handler;
        private final AdminApiDriver<String, Long> driver;
        private final AdminApiFuture.SimpleAdminApiFuture<String, Long> future;

        public TestContext(
            Map<String, Integer> staticKeys,
            Map<String, String> dynamicKeys
        ) {
            Map<String, MockRequestScope> lookupScopes = new HashMap<>();
            staticKeys.forEach((key, brokerId) -> {
                MockRequestScope scope = new MockRequestScope(OptionalInt.of(brokerId), null);
                lookupScopes.put(key, scope);
            });

            dynamicKeys.forEach((key, context) -> {
                MockRequestScope scope = new MockRequestScope(OptionalInt.empty(), context);
                lookupScopes.put(key, scope);
            });

            MockLookupStrategy<String> lookupStrategy = new MockLookupStrategy<>(lookupScopes);
            this.handler = new MockAdminApiHandler<>(lookupStrategy);
            this.future = AdminApiFuture.forKeys(lookupStrategy.lookupScopes.keySet());

            this.driver = new AdminApiDriver<>(
                handler,
                future,
                time.milliseconds() + API_TIMEOUT_MS,
                RETRY_BACKOFF_MS,
                RETRY_BACKOFF_MAX_MS,
                new LogContext()
            );

            staticKeys.forEach((key, brokerId) ->
                assertMappedKey(this, key, brokerId)
            );

            dynamicKeys.keySet().forEach(key ->
                assertUnmappedKey(this, key)
            );
        }

        public static TestContext staticMapped(Map<String, Integer> staticKeys) {
            return new TestContext(staticKeys, Collections.emptyMap());
        }

        public static TestContext dynamicMapped(Map<String, String> dynamicKeys) {
            return new TestContext(Collections.emptyMap(), dynamicKeys);
        }

        private void assertLookupResponse(
            RequestSpec<String> requestSpec,
            LookupResult<String> result
        ) {
            requestSpec.keys.forEach(key ->
                assertUnmappedKey(this, key)
            );

            // The response is just a placeholder. The result is all we are interested in
            MetadataResponse response = new MetadataResponse(new MetadataResponseData(),
                ApiKeys.METADATA.latestVersion());
            driver.onResponse(time.milliseconds(), requestSpec, response, Node.noNode());

            result.mappedKeys.forEach((key, brokerId) ->
                assertMappedKey(this, key, brokerId)
            );

            result.failedKeys.forEach((key, exception) ->
                assertFailedKey(this, key, exception)
            );
        }

        private void assertResponse(
            RequestSpec<String> requestSpec,
            ApiResult<String, Long> result,
            Node node
        ) {
            int brokerId = requestSpec.scope.destinationBrokerId().orElseThrow(() ->
                new AssertionError("Fulfillment requests must specify a target brokerId"));

            requestSpec.keys.forEach(key ->
                assertMappedKey(this, key, brokerId)
            );

            // The response is just a placeholder. The result is all we are interested in
            MetadataResponse response = new MetadataResponse(new MetadataResponseData(),
                ApiKeys.METADATA.latestVersion());

            driver.onResponse(time.milliseconds(), requestSpec, response, node);

            result.unmappedKeys.forEach(key ->
                assertUnmappedKey(this, key)
            );

            result.failedKeys.forEach((key, exception) ->
                assertFailedKey(this, key, exception)
            );

            result.completedKeys.forEach((key, value) ->
                assertCompletedKey(this, key, value)
            );
        }

        private MockLookupStrategy<String> lookupStrategy() {
            return handler.lookupStrategy;
        }

        public void poll(
            Map<Set<String>, LookupResult<String>> expectedLookups,
            Map<Set<String>, ApiResult<String, Long>> expectedRequests
        ) {
            if (!expectedLookups.isEmpty()) {
                MockLookupStrategy<String> lookupStrategy = lookupStrategy();
                lookupStrategy.reset();
                expectedLookups.forEach(lookupStrategy::expectLookup);
            }

            handler.reset();
            expectedRequests.forEach(handler::expectRequest);

            List<RequestSpec<String>> requestSpecs = driver.poll();
            assertEquals(expectedLookups.size() + expectedRequests.size(), requestSpecs.size(),
                "Driver generated an unexpected number of requests");

            for (RequestSpec<String> requestSpec : requestSpecs) {
                Set<String> keys = requestSpec.keys;
                if (expectedLookups.containsKey(keys)) {
                    LookupResult<String> result = expectedLookups.get(keys);
                    assertLookupResponse(requestSpec, result);
                } else if (expectedRequests.containsKey(keys)) {
                    ApiResult<String, Long> result = expectedRequests.get(keys);
                    assertResponse(requestSpec, result, Node.noNode());
                } else {
                    fail("Unexpected request for keys " + keys);
                }
            }
        }
    }

    private static class MockLookupStrategy<K> implements AdminApiLookupStrategy<K> {
        private final Map<Set<K>, LookupResult<K>> expectedLookups = new HashMap<>();
        private final Map<K, MockRequestScope> lookupScopes;

        private MockLookupStrategy(Map<K, MockRequestScope> lookupScopes) {
            this.lookupScopes = lookupScopes;
        }

        @Override
        public ApiRequestScope lookupScope(K key) {
            return lookupScopes.get(key);
        }

        public void expectLookup(Set<K> keys, LookupResult<K> result) {
            expectedLookups.put(keys, result);
        }

        @Override
        public AbstractRequest.Builder<?> buildRequest(Set<K> keys) {
            // The request is just a placeholder in these tests
            assertTrue(expectedLookups.containsKey(keys), "Unexpected lookup request for keys " + keys);
            return new MetadataRequest.Builder(Collections.emptyList(), false);
        }

        @Override
        public LookupResult<K> handleResponse(Set<K> keys, AbstractResponse response) {
            return Optional.ofNullable(expectedLookups.get(keys)).orElseThrow(() ->
                new AssertionError("Unexpected fulfillment request for keys " + keys)
            );
        }

        public void reset() {
            expectedLookups.clear();
        }
    }

    private static class MockAdminApiHandler<K, V> extends AdminApiHandler.Batched<K, V> {
        private final Map<Set<K>, ApiResult<K, V>> expectedRequests = new HashMap<>();
        private final MockLookupStrategy<K> lookupStrategy;
        private final Map<K, Boolean> retriableUnsupportedVersionKeys;

        private MockAdminApiHandler(MockLookupStrategy<K> lookupStrategy) {
            this.lookupStrategy = lookupStrategy;
            this.retriableUnsupportedVersionKeys = new ConcurrentHashMap<>();
        }

        @Override
        public String apiName() {
            return "mock-api";
        }

        @Override
        public AdminApiLookupStrategy<K> lookupStrategy() {
            return lookupStrategy;
        }

        public void expectRequest(Set<K> keys, ApiResult<K, V> result) {
            expectedRequests.put(keys, result);
        }

        @Override
        public AbstractRequest.Builder<?> buildBatchedRequest(int brokerId, Set<K> keys) {
            // The request is just a placeholder in these tests
            assertTrue(expectedRequests.containsKey(keys), "Unexpected fulfillment request for keys " + keys);
            return new MetadataRequest.Builder(Collections.emptyList(), false);
        }

        @Override
        public ApiResult<K, V> handleResponse(Node broker, Set<K> keys, AbstractResponse response) {
            return Optional.ofNullable(expectedRequests.get(keys)).orElseThrow(() ->
                new AssertionError("Unexpected fulfillment request for keys " + keys)
            );
        }

        @Override
        public Map<K, Throwable> handleUnsupportedVersionException(
            int brokerId,
            UnsupportedVersionException exception,
            Set<K> keys
        ) {
            return keys
                .stream()
                .filter(k -> !retriableUnsupportedVersionKeys.containsKey(k))
                .collect(Collectors.toMap(k -> k, k -> exception));
        }

        public void reset() {
            expectedRequests.clear();
        }

        public void addRetriableUnsupportedVersionKey(K key) {
            retriableUnsupportedVersionKeys.put(key, Boolean.TRUE);
        }
    }

    private static <K, V> Map<K, V> map(K key, V value) {
        return Collections.singletonMap(key, value);
    }

    private static <K, V> Map<K, V> map(K k1, V v1, K k2, V v2) {
        HashMap<K, V> map = new HashMap<>(2);
        map.put(k1, v1);
        map.put(k2, v2);
        return map;
    }

    private static <K, V> Map<K, V> map(K k1, V v1, K k2, V v2, K k3, V v3) {
        HashMap<K, V> map = new HashMap<>(3);
        map.put(k1, v1);
        map.put(k2, v2);
        map.put(k3, v3);
        return map;
    }

    private static ApiResult<String, Long> completed(String key, Long value) {
        return new ApiResult<>(map(key, value), emptyMap(), Collections.emptyList());
    }

    private static ApiResult<String, Long> failed(String key, Throwable exception) {
        return new ApiResult<>(emptyMap(), map(key, exception), Collections.emptyList());
    }

    private static ApiResult<String, Long> unmapped(String... keys) {
        return new ApiResult<>(emptyMap(), emptyMap(), Arrays.asList(keys));
    }

    private static ApiResult<String, Long> completed(String k1, Long v1, String k2, Long v2) {
        return new ApiResult<>(map(k1, v1, k2, v2), emptyMap(), Collections.emptyList());
    }

    private static ApiResult<String, Long> emptyFulfillment() {
        return new ApiResult<>(emptyMap(), emptyMap(), Collections.emptyList());
    }

    private static LookupResult<String> failedLookup(String key, Throwable exception) {
        return new LookupResult<>(map(key, exception), emptyMap());
    }

    private static LookupResult<String> emptyLookup() {
        return new LookupResult<>(emptyMap(), emptyMap());
    }

    private static LookupResult<String> mapped(String key, Integer brokerId) {
        return new LookupResult<>(emptyMap(), map(key, brokerId));
    }

    private static LookupResult<String> mapped(String k1, Integer broker1, String k2, Integer broker2) {
        return new LookupResult<>(emptyMap(), map(k1, broker1, k2, broker2));
    }

}
