Repository: apache/kafka
File: share-coordinator/src/test/java/org/apache/kafka/coordinator/share/ShareCoordinatorShardTest.java (Java)

Code before:
        ).value().message()), shard.getShareStateMapValue(shareCoordinatorKey));
        assertEquals(0, shard.getLeaderMapValue(shareCoordinatorKey));
        verify(shard.getMetricsShard()).record(ShareCoordinatorMetrics.SHARE_COORDINATOR_WRITE_SENSOR_NAME);

        // State epoch stays same so share update.
        request = new WriteShareGroupStateRequestData()
            .setGroupId(GROUP_ID)
            .setTopics(List.of(new WriteShareGroupStateRequestData.WriteStateData()
                .setTopicId(TOPIC_ID)
                .setPartitions(List.of(new WriteShareGroupStateRequestData.PartitionData()
                    .setPartition(PARTITION)
                    .setStartOffset(0)
                    .setStateEpoch(stateEpoch)
                    .setLeaderEpoch(0)
                    .setStateBatches(List.of(new WriteShareGroupStateRequestData.StateBatch()
                        .setFirstOffset(0)
                        .setLastOffset(10)
                        .setDeliveryCount((short) 2)
                        .setDeliveryState((byte) 0)))))));

        result = shard.writeState(request);

        shard.replay(0L, 0L, (short) 0, result.records().get(0));

        expectedData = WriteShareGroupStateResponse.toResponseData(TOPIC_ID, PARTITION);
        expectedRecords = List.of(ShareCoordinatorRecordHelpers.newShareSnapshotUpdateRecord(
            GROUP_ID, TOPIC_ID, PARTITION, ShareGroupOffset.fromRequest(request.topics().get(0).partitions().get(0), TIME.milliseconds())
        ));

        // Snapshot epoch did not increase
        assertEquals(0, shard.getShareStateMapValue(shareCoordinatorKey).snapshotEpoch());
        assertEquals(expectedData, result.response());
        assertEquals(expectedRecords, result.records());

        assertEquals(groupOffset(ShareCoordinatorRecordHelpers.newShareSnapshotRecord(
            GROUP_ID, TOPIC_ID, PARTITION, ShareGroupOffset.fromRequest(request.topics().get(0).partitions().get(0), TIME.milliseconds())
        ).value().message()), shard.getShareStateMapValue(shareCoordinatorKey));
        assertEquals(0, shard.getLeaderMapValue(shareCoordinatorKey));
        verify(shard.getMetricsShard(), times(2)).record(ShareCoordinatorMetrics.SHARE_COORDINATOR_WRITE_SENSOR_NAME);

        // State epoch incremented so share snapshot.
        request = new WriteShareGroupStateRequestData()
            .setGroupId(GROUP_ID)
            .setTopics(List.of(new WriteShareGroupStateRequestData.WriteStateData()
                .setTopicId(TOPIC_ID)
                .setPartitions(List.of(new WriteShareGroupStateRequestData.PartitionData()
                    .setPartition(PARTITION)
                    .setStartOffset(0)
                    .setStateEpoch(stateEpoch + 1)   // incremented
                    .setLeaderEpoch(0)
                    .setStateBatches(List.of(new WriteShareGroupStateRequestData.StateBatch()

Review comment (refactor):
It would be better if the method `ShareCoordinatorRecordHelpers.newShareSnapshotUpdateRecord` was actually called `newShareUpdateRecord`. There are only 4 instances so I think it's worth making this change too.

Code after:
        ).value().message()), shard.getShareStateMapValue(shareCoordinatorKey));
        assertEquals(0, shard.getLeaderMapValue(shareCoordinatorKey));
        verify(shard.getMetricsShard()).record(ShareCoordinatorMetrics.SHARE_COORDINATOR_WRITE_SENSOR_NAME);

        // State epoch stays same so share update.
        request = new WriteShareGroupStateRequestData()
            .setGroupId(GROUP_ID)
            .setTopics(List.of(new WriteShareGroupStateRequestData.WriteStateData()
                .setTopicId(TOPIC_ID)
                .setPartitions(List.of(new WriteShareGroupStateRequestData.PartitionData()
                    .setPartition(PARTITION)
                    .setStartOffset(0)
                    .setStateEpoch(stateEpoch)
                    .setLeaderEpoch(0)
                    .setStateBatches(List.of(new WriteShareGroupStateRequestData.StateBatch()
                        .setFirstOffset(0)
                        .setLastOffset(10)
                        .setDeliveryCount((short) 2)
                        .setDeliveryState((byte) 0)))))));

        result = shard.writeState(request);

        shard.replay(0L, 0L, (short) 0, result.records().get(0));

        expectedData = WriteShareGroupStateResponse.toResponseData(TOPIC_ID, PARTITION);
        expectedRecords = List.of(ShareCoordinatorRecordHelpers.newShareUpdateRecord(
            GROUP_ID, TOPIC_ID, PARTITION, ShareGroupOffset.fromRequest(request.topics().get(0).partitions().get(0), TIME.milliseconds())
        ));

        // Snapshot epoch did not increase
        assertEquals(0, shard.getShareStateMapValue(shareCoordinatorKey).snapshotEpoch());
        assertEquals(expectedData, result.response());
        assertEquals(expectedRecords, result.records());

        assertEquals(groupOffset(ShareCoordinatorRecordHelpers.newShareSnapshotRecord(
            GROUP_ID, TOPIC_ID, PARTITION, ShareGroupOffset.fromRequest(request.topics().get(0).partitions().get(0), TIME.milliseconds())
        ).value().message()), shard.getShareStateMapValue(shareCoordinatorKey));
        assertEquals(0, shard.getLeaderMapValue(shareCoordinatorKey));
        verify(shard.getMetricsShard(), times(2)).record(ShareCoordinatorMetrics.SHARE_COORDINATOR_WRITE_SENSOR_NAME);

        // State epoch incremented so share snapshot.
        request = new WriteShareGroupStateRequestData()
            .setGroupId(GROUP_ID)
            .setTopics(List.of(new WriteShareGroupStateRequestData.WriteStateData()
                .setTopicId(TOPIC_ID)
                .setPartitions(List.of(new WriteShareGroupStateRequestData.PartitionData()
                    .setPartition(PARTITION)
                    .setStartOffset(0)
                    .setStateEpoch(stateEpoch + 1)   // incremented
                    .setLeaderEpoch(0)
                    .setStateBatches(List.of(new WriteShareGroupStateRequestData.StateBatch()