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

Code before:
  }

  /**
   * If a partition becomes a follower and the leader is unchanged it should check for truncation
   * if the epoch has increased by more than one (which suggests it has missed an update). For
   * IBP version 2.7 onwards, we don't require this since we can truncate at any time based
   * on diverging epochs returned in fetch responses.
   */
  private def verifyBecomeFollowerWhenLeaderIsUnchangedButMissedLeaderUpdate(extraProps: Properties,
                                                                             expectTruncation: Boolean): Unit = {
    val followerBrokerId = 0
    val leaderBrokerId = 1
    var leaderEpoch = 1
    val leaderEpochIncrement = 2
    val countDownLatch = new CountDownLatch(1)
    val offsetFromLeader = 5
    // Prepare the mocked components for the test
    val (replicaManager, mockLogMgr) = prepareReplicaManagerAndLogManager(new MockTimer(time),
      0, leaderEpoch + leaderEpochIncrement, followerBrokerId, leaderBrokerId, countDownLatch,
      expectTruncation = expectTruncation, localLogOffset = Optional.of(10), offsetFromLeader = offsetFromLeader, extraProps = extraProps, topicId = Optional.of(topicId))

    try {
      // Initialize partition state to follower, with leader = 1, leaderEpoch = 1
      val directoryIds = replicaManager.logManager.directoryIdsSet.toList
      val initialFollowerDelta = topicsCreateDelta( startId = followerBrokerId, isStartIdLeader = false, partitions = List(0), directoryIds = directoryIds, topic, topicIds(topic), leaderEpoch)
      val initialImage = imageFromTopics(initialFollowerDelta.apply())
      replicaManager.applyDelta(initialFollowerDelta, initialImage)

      // Verify log created and partition is hosted
      val localLog = replicaManager.localLog(topicPartition)
      assertTrue(localLog.isDefined, "Log should be created for follower after applyDelta")
      val hostedPartition = replicaManager.getPartition(topicPartition)
      assertTrue(hostedPartition.isInstanceOf[HostedPartition.Online])

      // Make local partition a follower - because epoch increased by more than 1, truncation should
      // trigger even though leader does not change
      leaderEpoch += leaderEpochIncrement
      val epochJumpDelta = topicsCreateDelta( followerBrokerId, isStartIdLeader = false, partitions = List(0), directoryIds = directoryIds, topicName = topic, topicId = topicIds(topic), leaderEpoch = leaderEpoch)
      val newImage = imageFromTopics(epochJumpDelta.apply())
      replicaManager.applyDelta(epochJumpDelta, newImage)

      assertTrue(countDownLatch.await(1000L, TimeUnit.MILLISECONDS))

      // Truncation should have happened once
      if (expectTruncation) {
        verify(mockLogMgr).truncateTo(Map(topicPartition -> offsetFromLeader), isFuture = false)
      }

      verify(mockLogMgr).finishedInitializingLog(ArgumentMatchers.eq(topicPartition), any())
    } finally {
      replicaManager.shutdown(checkpointHW = false)

Review comment (suggestion):
This should fix the test. You can update based on this.
```suggestion
      val tp = new TopicPartition(topic, topicPartition)
      val partition = replicaManager.createPartition(tp)
      val offsetCheckpoints = new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints.asJava)
      partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints, None)

      val initialFollowerDelta = topicsCreateDelta(startId = followerBrokerId, isStartIdLeader = false, partitions = List(0), List.empty, topic, topicIds(topic), leaderEpoch)
      val initialImage = imageFromTopics(initialFollowerDelta.apply())
      replicaManager.applyDelta(initialFollowerDelta, initialImage)
```

Code after:
    // Prepare the mocked components for the test
    val (replicaManager, mockLogMgr) = prepareReplicaManagerAndLogManager(new MockTimer(time),
      topicPartition.partition(), leaderEpoch + leaderEpochIncrement, followerBrokerId, leaderBrokerId, countDownLatch,
      expectTruncation = false, localLogOffset = Optional.of(10), offsetFromLeader = offsetFromLeader, extraProps = extraProps, topicId = Optional.of(topicId))

    try {
      // Initialize partition state to follower, with leader = 1, leaderEpoch = 1
      val partition = replicaManager.createPartition(topicPartition)
      val offsetCheckpoints = new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints.asJava)
      partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints, None)
      val followerDelta = topicsCreateDelta(startId = followerBrokerId, isStartIdLeader = false, partitions = List(topicPartition.partition()), List.empty, topic, topicIds(topic), leaderEpoch)
      replicaManager.applyDelta(followerDelta, imageFromTopics(followerDelta.apply()))

      // Verify log created and partition is hosted
      val localLog = replicaManager.localLog(topicPartition)
      assertTrue(localLog.isDefined, "Log should be created for follower after applyDelta")
      val hostedPartition = replicaManager.getPartition(topicPartition)
      assertTrue(hostedPartition.isInstanceOf[HostedPartition.Online])

      // Make local partition a follower - because epoch increased by more than 1, truncation should
      // trigger even though leader does not change
      leaderEpoch += leaderEpochIncrement
      val epochJumpDelta = topicsCreateDelta(startId = followerBrokerId, isStartIdLeader = false, partitions = List(topicPartition.partition()), List.empty, topic, topicIds(topic), leaderEpoch)
      replicaManager.applyDelta(epochJumpDelta, imageFromTopics(epochJumpDelta.apply()))

      assertTrue(countDownLatch.await(1000L, TimeUnit.MILLISECONDS))


      verify(mockLogMgr).finishedInitializingLog(ArgumentMatchers.eq(topicPartition), any())
    } finally {
      replicaManager.shutdown(checkpointHW = false)
    }
  }

  @Test
  def testReplicaSelector(): Unit = {
    val topicPartition = 0
    val followerBrokerId = 0
    val leaderBrokerId = 1
    val leaderEpoch = 1
    val leaderEpochIncrement = 2
    val aliveBrokerIds = Seq[Integer](followerBrokerId, leaderBrokerId)
    val countDownLatch = new CountDownLatch(1)

    // Prepare the mocked components for the test
    val (replicaManager, _) = prepareReplicaManagerAndLogManager(new MockTimer(time),
      topicPartition, leaderEpoch + leaderEpochIncrement, followerBrokerId,
      leaderBrokerId, countDownLatch, expectTruncation = true)

    try {
      val tp = new TopicPartition(topic, topicPartition)