ControllerResult<AssignReplicasToDirsResponseData> handleAssignReplicasToDirs(AssignReplicasToDirsRequestData request) {
        if (!featureControl.metadataVersion().isDirectoryAssignmentSupported()) {
            throw new UnsupportedVersionException("Directory assignment is not supported yet.");
        }
        int brokerId = request.brokerId();
        clusterControl.checkBrokerEpoch(brokerId, request.brokerEpoch());
        BrokerRegistration brokerRegistration = clusterControl.brokerRegistrations().get(brokerId);
        if (brokerRegistration == null) {
            throw new BrokerIdNotRegisteredException("Broker ID " + brokerId + " is not currently registered");
        }
        List<ApiMessageAndVersion> records = new ArrayList<>();
        AssignReplicasToDirsResponseData response = new AssignReplicasToDirsResponseData();
        Set<TopicIdPartition> leaderAndIsrUpdates = new HashSet<>();
        for (AssignReplicasToDirsRequestData.DirectoryData reqDir : request.directories()) {
            Uuid dirId = reqDir.id();
            boolean directoryIsOffline = !brokerRegistration.hasOnlineDir(dirId);
            AssignReplicasToDirsResponseData.DirectoryData resDir = new AssignReplicasToDirsResponseData.DirectoryData().setId(dirId);
            for (AssignReplicasToDirsRequestData.TopicData reqTopic : reqDir.topics()) {
                Uuid topicId = reqTopic.topicId();
                Errors topicError = Errors.NONE;
                TopicControlInfo topicInfo = this.topics.get(topicId);
                if (topicInfo == null) {
                    log.warn("AssignReplicasToDirsRequest from broker {} references unknown topic ID {}", brokerId, topicId);
                    topicError = Errors.UNKNOWN_TOPIC_ID;
                }
                AssignReplicasToDirsResponseData.TopicData resTopic = new AssignReplicasToDirsResponseData.TopicData().setTopicId(topicId);
                for (AssignReplicasToDirsRequestData.PartitionData reqPartition : reqTopic.partitions()) {
                    int partitionIndex = reqPartition.partitionIndex();
                    Errors partitionError = topicError;
                    if (topicError == Errors.NONE) {
                        String topicName = topicInfo.name;
                        PartitionRegistration partitionRegistration = topicInfo.parts.get(partitionIndex);
                        if (partitionRegistration == null) {
                            log.warn("AssignReplicasToDirsRequest from broker {} references unknown partition {}-{}", brokerId, topicName, partitionIndex);
                            partitionError = Errors.UNKNOWN_TOPIC_OR_PARTITION;
                        } else if (!Replicas.contains(partitionRegistration.replicas, brokerId)) {
                            log.warn("AssignReplicasToDirsRequest from broker {} references non assigned partition {}-{}", brokerId, topicName, partitionIndex);
                            partitionError = Errors.NOT_LEADER_OR_FOLLOWER;
                        } else {
                            Optional<ApiMessageAndVersion> partitionChangeRecord = new PartitionChangeBuilder(
                                    partitionRegistration,
                                    topicId,
                                    partitionIndex,
                                    new LeaderAcceptor(clusterControl, partitionRegistration),
                                    featureControl.metadataVersion(),
                                    getTopicEffectiveMinIsr(topicName)
                            )
                                    .setDirectory(brokerId, dirId)
                                    .setDefaultDirProvider(clusterDescriber)
                                    .build();
                            partitionChangeRecord.ifPresent(records::add);
                            if (directoryIsOffline) {
                                leaderAndIsrUpdates.add(new TopicIdPartition(topicId, partitionIndex));
                            }
                            if (log.isDebugEnabled()) {
                                log.debug("Broker {} assigned partition {}:{} to {} dir {}",
                                    brokerId, topics.get(topicId).name(), partitionIndex,
                                    directoryIsOffline ? "OFFLINE" : "ONLINE", dirId);
                            }
                        }
                    }
                    resTopic.partitions().add(new AssignReplicasToDirsResponseData.PartitionData().
                            setPartitionIndex(partitionIndex).
                            setErrorCode(partitionError.code()));
                }
                resDir.topics().add(resTopic);
            }
            response.directories().add(resDir);
        }
        if (!leaderAndIsrUpdates.isEmpty()) {
            generateLeaderAndIsrUpdates("offline-dir-assignment", brokerId, NO_LEADER, NO_LEADER, records, leaderAndIsrUpdates.iterator());
        }
        return ControllerResult.of(records, response);
    }
@Test
    void testHandleAssignReplicasToDirs() {
        ReplicationControlTestContext ctx = new ReplicationControlTestContext.Builder().build();
        Uuid dir1b1 = Uuid.fromString("hO2YI5bgRUmByNPHiHxjNQ");
        Uuid dir2b1 = Uuid.fromString("R3Gb1HLoTzuKMgAkH5Vtpw");
        Uuid dir1b2 = Uuid.fromString("TBGa8UayQi6KguqF5nC0sw");
        Uuid offlineDir = Uuid.fromString("zvAf9BKZRyyrEWz4FX2nLA");
        ctx.registerBrokersWithDirs(1, asList(dir1b1, dir2b1), 2, singletonList(dir1b2));
        ctx.unfenceBrokers(1, 2);
        Uuid topicA = ctx.createTestTopic("a", new int[][]{new int[]{1, 2}, new int[]{1, 2}, new int[]{1, 2}}).topicId();
        Uuid topicB = ctx.createTestTopic("b", new int[][]{new int[]{1, 2}, new int[]{1, 2}}).topicId();
        Uuid topicC = ctx.createTestTopic("c", new int[][]{new int[]{2}}).topicId();

        ControllerResult<AssignReplicasToDirsResponseData> controllerResult = ctx.assignReplicasToDirs(1, new HashMap<TopicIdPartition, Uuid>() {{
                put(new TopicIdPartition(topicA, 0), dir1b1);
                put(new TopicIdPartition(topicA, 1), dir2b1);
                put(new TopicIdPartition(topicA, 2), offlineDir); // unknown/offline dir
                put(new TopicIdPartition(topicB, 0), dir1b1);
                put(new TopicIdPartition(topicB, 1), DirectoryId.LOST);
                put(new TopicIdPartition(Uuid.fromString("nLU9hKNXSZuMe5PO2A4dVQ"), 1), dir2b1); // expect UNKNOWN_TOPIC_ID
                put(new TopicIdPartition(topicA, 137), dir1b1); // expect UNKNOWN_TOPIC_OR_PARTITION
                put(new TopicIdPartition(topicC, 0), dir1b1); // expect NOT_LEADER_OR_FOLLOWER
            }});

        assertEquals(AssignmentsHelper.normalize(AssignmentsHelper.buildResponseData((short) 0, 0, new HashMap<Uuid, Map<TopicIdPartition, Errors>>() {{
                put(dir1b1, new HashMap<TopicIdPartition, Errors>() {{
                        put(new TopicIdPartition(topicA, 0), NONE);
                        put(new TopicIdPartition(topicA, 137), UNKNOWN_TOPIC_OR_PARTITION);
                        put(new TopicIdPartition(topicB, 0), NONE);
                        put(new TopicIdPartition(topicC, 0), NOT_LEADER_OR_FOLLOWER);
                    }});
                put(dir2b1, new HashMap<TopicIdPartition, Errors>() {{
                        put(new TopicIdPartition(topicA, 1), NONE);
                        put(new TopicIdPartition(Uuid.fromString("nLU9hKNXSZuMe5PO2A4dVQ"), 1), UNKNOWN_TOPIC_ID);
                    }});
                put(offlineDir, new HashMap<TopicIdPartition, Errors>() {{
                        put(new TopicIdPartition(topicA, 2), NONE);
                    }});
                put(DirectoryId.LOST, new HashMap<TopicIdPartition, Errors>() {{
                        put(new TopicIdPartition(topicB, 1), NONE);
                    }});
            }})), AssignmentsHelper.normalize(controllerResult.response()));
        short recordVersion = ctx.featureControl.metadataVersion().partitionChangeRecordVersion();
        assertEquals(sortPartitionChangeRecords(asList(
                new ApiMessageAndVersion(
                        new PartitionChangeRecord().setTopicId(topicA).setPartitionId(0)
                                .setDirectories(asList(dir1b1, dir1b2)), recordVersion),
                new ApiMessageAndVersion(
                        new PartitionChangeRecord().setTopicId(topicA).setPartitionId(1).
                                setDirectories(asList(dir2b1, dir1b2)), recordVersion),
                new ApiMessageAndVersion(
                        new PartitionChangeRecord().setTopicId(topicA).setPartitionId(2).
                                setDirectories(asList(offlineDir, dir1b2)), recordVersion),
                new ApiMessageAndVersion(
                        new PartitionChangeRecord().setTopicId(topicB).setPartitionId(0).
                                setDirectories(asList(dir1b1, dir1b2)), recordVersion),
                new ApiMessageAndVersion(
                        new PartitionChangeRecord().setTopicId(topicB).setPartitionId(1).
                                setDirectories(asList(DirectoryId.LOST, dir1b2)), recordVersion),

                // In addition to the directory assignment changes we expect two additional records,
                // which elect new leaders for:
                //   - a-2 which has been assigned to a directory which is not an online directory (unknown/offline)
                //   - b-1 which has been assigned to an offline directory.
                new ApiMessageAndVersion(
                        new PartitionChangeRecord().setTopicId(topicA).setPartitionId(2).
                                setIsr(singletonList(2)).setLeader(2), recordVersion),
                new ApiMessageAndVersion(
                        new PartitionChangeRecord().setTopicId(topicB).setPartitionId(1).
                                setIsr(singletonList(2)).setLeader(2), recordVersion)
        )), sortPartitionChangeRecords(controllerResult.records()));

        ctx.replay(controllerResult.records());
        assertEquals(new HashSet<TopicIdPartition>() {{
                add(new TopicIdPartition(topicA, 0));
                add(new TopicIdPartition(topicA, 1));
                add(new TopicIdPartition(topicB, 0));
            }}, RecordTestUtils.iteratorToSet(ctx.replicationControl.brokersToIsrs().iterator(1, true)));
        assertEquals(new HashSet<TopicIdPartition>() {{
                add(new TopicIdPartition(topicA, 2));
                add(new TopicIdPartition(topicB, 1));
                add(new TopicIdPartition(topicC, 0));
            }},
            RecordTestUtils.iteratorToSet(ctx.replicationControl.brokersToIsrs().iterator(2, true)));
    }