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

Code before:
      waitUntilTrue(() => brokerTopicStats.topicStats(tp0.topic()).buildRemoteLogAuxStateRequestRate.count > 0,
        "Should have buildRemoteLogAuxStateRequestRate count > 0, but got:" + brokerTopicStats.topicStats(tp0.topic()).buildRemoteLogAuxStateRequestRate.count)
      assertEquals(0, brokerTopicStats.topicStats(tp0.topic()).failedBuildRemoteLogAuxStateRate.count)
      // Verify aggregate metrics
      waitUntilTrue(() => brokerTopicStats.allTopicsStats.buildRemoteLogAuxStateRequestRate.count > 0,
        "Should have all topic buildRemoteLogAuxStateRequestRate count > 0, but got:" + brokerTopicStats.allTopicsStats.buildRemoteLogAuxStateRequestRate.count)
      assertEquals(0, brokerTopicStats.allTopicsStats.failedBuildRemoteLogAuxStateRate.count)
    } finally {
      replicaManager.shutdown(checkpointHW = false)
    }
  }

  @Test
  def testFailedBuildRemoteLogAuxStateMetrics(): Unit = {
    val tp0 = new TopicPartition(topic, 0)

    val remoteLogManager = mock(classOf[RemoteLogManager])
    val storageManager = mock(classOf[RemoteStorageManager])
    when(storageManager.fetchIndex(any(), any())).thenReturn(new ByteArrayInputStream("0".getBytes()))
    when(remoteLogManager.storageManager()).thenReturn(storageManager)

    val replicaManager = setupReplicaManagerWithMockedPurgatories(new MockTimer(time), aliveBrokerIds = Seq(0, 1, 2), enableRemoteStorage = true, shouldMockLog = true, remoteLogManager = Some(remoteLogManager), buildRemoteLogAuxState = true)
    try {
      val offsetCheckpoints = new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints.asJava)
      replicaManager.createPartition(tp0).createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints, None)
      val topicIds = Map(tp0.topic -> topicId).asJava

      // Verify the metrics for build remote log state and for failures is zero before replicas start to fetch
      assertEquals(0, brokerTopicStats.topicStats(tp0.topic()).buildRemoteLogAuxStateRequestRate.count)
      assertEquals(0, brokerTopicStats.topicStats(tp0.topic()).failedBuildRemoteLogAuxStateRate.count)
      // Verify aggregate metrics
      assertEquals(0, brokerTopicStats.allTopicsStats.buildRemoteLogAuxStateRequestRate.count)
      assertEquals(0, brokerTopicStats.allTopicsStats.failedBuildRemoteLogAuxStateRate.count)

      val delta = new TopicsDelta(TopicsImage.EMPTY)
      delta.replay(new TopicRecord().setName(topic).setTopicId(topicIds.get(topic)))
      val record = new PartitionRecord()
        .setPartitionId(0)
        .setTopicId(topicIds.get(topic))
        .setReplicas(util.List.of(0, 1))
        .setIsr(util.List.of(0, 1))
        .setRemovingReplicas(util.List.of())
        .setAddingReplicas(util.List.of())
        .setLeader(1)
        .setLeaderEpoch(0)
        .setPartitionEpoch(0)
      delta.replay(record)
      val leaderMetadataImage = imageFromTopics(delta.apply())
      replicaManager.applyDelta(delta, leaderMetadataImage)

      // Replicas fetch from the leader periodically, therefore we check that the metric value is increasing

Review comment (refactor):
Could you please remove this local variable and use class variable `topicIds` instead?

Code after:
      val partition0Replicas = Seq[Integer](0, 1).asJava

      // Verify the metrics for build remote log state and for failures is zero before replicas start to fetch
      assertEquals(0, brokerTopicStats.topicStats(tp0.topic()).buildRemoteLogAuxStateRequestRate.count)
      assertEquals(0, brokerTopicStats.topicStats(tp0.topic()).failedBuildRemoteLogAuxStateRate.count)
      // Verify aggregate metrics
      assertEquals(0, brokerTopicStats.allTopicsStats.buildRemoteLogAuxStateRequestRate.count)
      assertEquals(0, brokerTopicStats.allTopicsStats.failedBuildRemoteLogAuxStateRate.count)

      val leaderDelta = createLeaderDelta(topicId, tp0, leaderId = 1, replicas = partition0Replicas, isr = partition0Replicas)
      val leaderMetadataImage = imageFromTopics(leaderDelta.apply())
      replicaManager.applyDelta(leaderDelta, leaderMetadataImage)

      // Replicas fetch from the leader periodically, therefore we check that the metric value is increasing
      // We expect failedBuildRemoteLogAuxStateRate to increase because fetchRemoteLogSegmentMetadata returns RemoteStorageException
      TestUtils.waitUntilTrue(() => brokerTopicStats.topicStats(tp0.topic()).buildRemoteLogAuxStateRequestRate.count > 0,
        "Should have buildRemoteLogAuxStateRequestRate count > 0, but got:" + brokerTopicStats.topicStats(tp0.topic()).buildRemoteLogAuxStateRequestRate.count)
      TestUtils.waitUntilTrue(() => brokerTopicStats.topicStats(tp0.topic()).failedBuildRemoteLogAuxStateRate.count > 0,
        "Should have failedBuildRemoteLogAuxStateRate count > 0, but got:" + brokerTopicStats.topicStats(tp0.topic()).failedBuildRemoteLogAuxStateRate.count)
      // Verify aggregate metrics
      TestUtils.waitUntilTrue(() => brokerTopicStats.allTopicsStats.buildRemoteLogAuxStateRequestRate.count > 0,
        "Should have all topic buildRemoteLogAuxStateRequestRate count > 0, but got:" + brokerTopicStats.allTopicsStats.buildRemoteLogAuxStateRequestRate.count)
      TestUtils.waitUntilTrue(() => brokerTopicStats.allTopicsStats.failedBuildRemoteLogAuxStateRate.count > 0,
        "Should have all topic failedBuildRemoteLogAuxStateRate count > 0, but got:" + brokerTopicStats.allTopicsStats.failedBuildRemoteLogAuxStateRate.count)
    } finally {
      replicaManager.shutdown(checkpointHW = false)
    }
  }

  private def setupMockLog(path: String): UnifiedLog = {
    val mockLog = mock(classOf[UnifiedLog])
    val partitionDir = new File(path, s"$topic-0")
    partitionDir.mkdir()
    when(mockLog.dir).thenReturn(partitionDir)
    when(mockLog.parentDir).thenReturn(path)
    when(mockLog.topicId).thenReturn(Optional.of(topicId))
    when(mockLog.topicPartition).thenReturn(new TopicPartition(topic, 0))
    when(mockLog.highWatermark).thenReturn(highHW)
    when(mockLog.updateHighWatermark(anyLong())).thenReturn(0L)
    when(mockLog.logEndOffsetMetadata).thenReturn(new LogOffsetMetadata(10))
    when(mockLog.maybeIncrementHighWatermark(any(classOf[LogOffsetMetadata]))).thenReturn(Optional.empty)
    when(mockLog.endOffsetForEpoch(anyInt())).thenReturn(Optional.empty)
    // try to return a high start offset to cause OffsetOutOfRangeException at the 1st time
    when(mockLog.logStartOffset).thenReturn(endOffset).thenReturn(startOffset)
    when(mockLog.logEndOffset).thenReturn(endOffset)
    when(mockLog.localLogStartOffset()).thenReturn(endOffset - 10)
    when(mockLog.leaderEpochCache).thenReturn(mock(classOf[LeaderEpochFileCache]))
    when(mockLog.latestEpoch).thenReturn(Optional.of(0))
    val producerStateManager = mock(classOf[ProducerStateManager])
    when(mockLog.producerStateManager).thenReturn(producerStateManager)