Repository: apache/kafka
File: core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala (Scala)

Code before:
      var fetchResult = fetchPartitionAsConsumer(replicaManager, tidp0, partitionData, clientMetadata = Some(clientMetadata))
      assertEquals(Errors.NONE, fetchResult.assertFired.error)

      partitionData = new FetchRequest.PartitionData(Uuid.ZERO_UUID, 0L, 0L, 100,
        Optional.empty())
      fetchResult = fetchPartitionAsConsumer(replicaManager, tidp0, partitionData, clientMetadata = Some(clientMetadata))
      assertEquals(Errors.NONE, fetchResult.assertFired.error)
    } finally {
      replicaManager.shutdown(checkpointHW = false)
    }
  }

  @Test
  def testVerificationForTransactionalPartitionsOnly(): Unit = {
    val tp0 = new TopicPartition(topic, 0)
    val tp1 = new TopicPartition(topic, 1)
    val producerId = 24L
    val producerEpoch = 0.toShort
    val sequence = 0
    val addPartitionsToTxnManager = mock(classOf[AddPartitionsToTxnManager])

    val replicaManager = setUpReplicaManagerWithMockedAddPartitionsToTxnManager(addPartitionsToTxnManager, List(tp0, tp1))
    try {
      // Step 1: Apply KRaft metadata to assign tp0 and tp1 as leaders
      val directoryIds = replicaManager.logManager.directoryIdsSet.toList
      val leaderDelta = topicsCreateDelta(startId = 0, isStartIdLeader = true, partitions = List(0, 1), directoryIds = directoryIds, topic, topicIds(topic))
      val leaderImage = imageFromTopics(leaderDelta.apply())
      replicaManager.applyDelta(leaderDelta, leaderImage)

      // If we supply no transactional ID and idempotent records, we do not verify.
      val idempotentRecords = MemoryRecords.withIdempotentRecords(Compression.NONE, producerId, producerEpoch, sequence,
        new SimpleRecord("message".getBytes))
      handleProduceAppend(replicaManager, tp0, idempotentRecords, transactionalId = null)
      verify(addPartitionsToTxnManager, times(0)).addOrVerifyTransaction(any(), any(), any(), any(), any[AddPartitionsToTxnManager.AppendCallback](), any())
      assertEquals(VerificationGuard.SENTINEL, getVerificationGuard(replicaManager, tp0, producerId))

      // If we supply a transactional ID and some transactional and some idempotent records, we should only verify the topic partition with transactional records.
      val transactionalRecords = MemoryRecords.withTransactionalRecords(Compression.NONE, producerId, producerEpoch, sequence + 1,
        new SimpleRecord("message".getBytes))

      val idempotentRecords2 = MemoryRecords.withIdempotentRecords(Compression.NONE, producerId, producerEpoch, sequence,
        new SimpleRecord("message".getBytes))
      handleProduceAppendToMultipleTopics(replicaManager, Map(tp0 -> transactionalRecords, tp1 -> idempotentRecords2), transactionalId)
      verify(addPartitionsToTxnManager, times(1)).addOrVerifyTransaction(
        ArgumentMatchers.eq(transactionalId),
        ArgumentMatchers.eq(producerId),
        ArgumentMatchers.eq(producerEpoch),
        ArgumentMatchers.eq(util.List.of(tp0)),
        any[AddPartitionsToTxnManager.AppendCallback](),
        any()
      )

Review comment (suggestion):
Fix the test:
```suggestion
      val leaderDelta = topicsCreateDelta(startId = 0, isStartIdLeader = false, partitions = List(0, 1), List.empty, topic, topicIds(topic))
```


Code after:
        new SimpleRecord("message".getBytes))

      val idempotentRecords2 = MemoryRecords.withIdempotentRecords(Compression.NONE, producerId, producerEpoch, sequence,
        new SimpleRecord("message".getBytes))
      handleProduceAppendToMultipleTopics(replicaManager, Map(tp0 -> transactionalRecords, tp1 -> idempotentRecords2), transactionalId)
      verify(addPartitionsToTxnManager, times(1)).addOrVerifyTransaction(
        ArgumentMatchers.eq(transactionalId),
        ArgumentMatchers.eq(producerId),
        ArgumentMatchers.eq(producerEpoch),
        ArgumentMatchers.eq(util.List.of(tp0)),
        any[AddPartitionsToTxnManager.AppendCallback](),
        any()
      )
      assertNotEquals(VerificationGuard.SENTINEL, getVerificationGuard(replicaManager, tp0, producerId))
      assertEquals(VerificationGuard.SENTINEL, getVerificationGuard(replicaManager, tp1, producerId))
    } finally {
      replicaManager.shutdown(checkpointHW = false)
    }
  }

  @ParameterizedTest
  @EnumSource(value = classOf[AppendOrigin], names = Array("CLIENT", "COORDINATOR"))
  def testTransactionVerificationFlow(appendOrigin: AppendOrigin): Unit = {
    val tp0 = new TopicPartition(topic, 0)
    val producerId = 24L
    val producerEpoch = 0.toShort
    val sequence = 6
    val addPartitionsToTxnManager = mock(classOf[AddPartitionsToTxnManager])
    val brokerList = Seq[Integer](0, 1).asJava

    val replicaManager = setUpReplicaManagerWithMockedAddPartitionsToTxnManager(addPartitionsToTxnManager, List(tp0))
    try {
      val leaderDelta = createLeaderDelta(topicId, tp0, leaderId = 1, replicas = brokerList, isr = brokerList)
      val leaderMetadataImage = imageFromTopics(leaderDelta.apply())
      replicaManager.applyDelta(leaderDelta, leaderMetadataImage)

      // Append some transactional records.
      val transactionalRecords = MemoryRecords.withTransactionalRecords(Compression.NONE, producerId, producerEpoch, sequence,
        new SimpleRecord("message".getBytes))

      // We should add these partitions to the manager to verify.
      val result = handleProduceAppend(replicaManager, tp0, transactionalRecords, origin = appendOrigin, transactionalId = transactionalId)
      val appendCallback = ArgumentCaptor.forClass(classOf[AddPartitionsToTxnManager.AppendCallback])
      verify(addPartitionsToTxnManager, times(1)).addOrVerifyTransaction(
        ArgumentMatchers.eq(transactionalId),
        ArgumentMatchers.eq(producerId),
        ArgumentMatchers.eq(producerEpoch),
        ArgumentMatchers.eq(util.List.of(tp0)),
        appendCallback.capture(),
        any()
      )