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

Code before:
      // Remove the first client-id quota, quotaTypesEnabled should remain unchanged
      clientQuotaManager.updateQuota(None, Some(ClientQuotaManager.ClientIdEntity("client1")), None)
      assertEquals(QuotaTypes.UserClientIdQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled, clientQuotaManager.quotaTypesEnabled)
      assertTrue(clientQuotaManager.quotasEnabled)

      // Remove the second client-id quota, quotaTypesEnabled should be QuotaTypes.UserClientIdQuotaEnabled
      clientQuotaManager.updateQuota(None, Some(ClientQuotaManager.ClientIdEntity("client2")), None)
      assertEquals(QuotaTypes.UserClientIdQuotaEnabled, clientQuotaManager.quotaTypesEnabled)
      assertTrue(clientQuotaManager.quotasEnabled)

      // Remove the first user-client-id quota, quotaTypesEnabled should be noQuotas as both user-client-id quotas has the same user client but different quota
      clientQuotaManager.updateQuota(Some(ClientQuotaManager.UserEntity("userA")), Some(ClientQuotaManager.ClientIdEntity("client1")), None)
      assertEquals(QuotaTypes.NoQuotas, clientQuotaManager.quotaTypesEnabled)
      assertFalse(clientQuotaManager.quotasEnabled)

      // Remove the second user-client-id quota, quotaTypesEnabled should be QuotaTypes.NoQuotas and quotasEnabled should be false
      clientQuotaManager.updateQuota(Some(ClientQuotaManager.UserEntity("userA")), Some(ClientQuotaManager.ClientIdEntity("client1")), None)
      assertEquals(QuotaTypes.NoQuotas, clientQuotaManager.quotaTypesEnabled)
      assertFalse(clientQuotaManager.quotasEnabled)
    } finally {
      clientQuotaManager.shutdown()
    }
  }

  @Test
  def testQuotaTypesEnabledUpdatesWithCustomCallback(): Unit = {
    val customQuotaCallback = new ClientQuotaCallback {
      val quotas = new HashMap[ClientQuotaEntity, Quota]()
      override def configure(configs: Map[String, _]): Unit = {}

      override def quotaMetricTags(quotaType: ClientQuotaType, principal: KafkaPrincipal, clientId: String): Map[String, String] = Collections.emptyMap()

      override def quotaLimit(quotaType: ClientQuotaType, metricTags: Map[String, String]): java.lang.Double = 1
      override def updateClusterMetadata(cluster: Cluster): Boolean = false

      override def updateQuota(quotaType: ClientQuotaType, entity: ClientQuotaEntity, newValue: Double): Unit = {
        quotas.put(entity.asInstanceOf[KafkaQuotaEntity], new Quota(newValue.toLong, true))
      }

      override def removeQuota(quotaType: ClientQuotaType, entity: ClientQuotaEntity): Unit = {
        quotas.remove(entity.asInstanceOf[KafkaQuotaEntity])
      }

      override def quotaResetRequired(quotaType: ClientQuotaType): Boolean = false

      override def getActiveQuotasEntities: Set[ClientQuotaEntity] =
        quotas.keySet().asInstanceOf[Set[ClientQuotaEntity]]

      override def close(): Unit = {}
    }
    val clientQuotaManager = new ClientQuotaManager(

Review comment (suggestion):
We don't rely on quotaTypesEnabled with custom callbacks, so don't need this test. Or change it to just verify that QuotaTypes.CustomQuotas is set.

Code after:
      // Remove the first client-id quota, quotaTypesEnabled should remain unchanged
      clientQuotaManager.updateQuota(None, Some(ClientQuotaManager.ClientIdEntity("client1")), None)
      assertEquals(QuotaTypes.UserClientIdQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled, clientQuotaManager.quotaTypesEnabled)
      assertTrue(clientQuotaManager.quotasEnabled)

      // Remove the second client-id quota, quotaTypesEnabled should be QuotaTypes.UserClientIdQuotaEnabled
      clientQuotaManager.updateQuota(None, Some(ClientQuotaManager.ClientIdEntity("client2")), None)
      assertEquals(QuotaTypes.UserClientIdQuotaEnabled, clientQuotaManager.quotaTypesEnabled)
      assertTrue(clientQuotaManager.quotasEnabled)

      // Remove the first user-client-id quota, quotaTypesEnabled should be noQuotas as both user-client-id quotas has the same user client but different quota
      clientQuotaManager.updateQuota(Some(ClientQuotaManager.UserEntity("userA")), Some(ClientQuotaManager.ClientIdEntity("client1")), None)
      assertEquals(QuotaTypes.NoQuotas, clientQuotaManager.quotaTypesEnabled)
      assertFalse(clientQuotaManager.quotasEnabled)

      // Remove the second user-client-id quota, quotaTypesEnabled should be QuotaTypes.NoQuotas and quotasEnabled should be false
      clientQuotaManager.updateQuota(Some(ClientQuotaManager.UserEntity("userA")), Some(ClientQuotaManager.ClientIdEntity("client1")), None)
      assertEquals(QuotaTypes.NoQuotas, clientQuotaManager.quotaTypesEnabled)
      assertFalse(clientQuotaManager.quotasEnabled)
    } finally {
      clientQuotaManager.shutdown()
    }
  }

  @Test
  def testQuotaTypesEnabledUpdatesWithCustomCallback(): Unit = {
    val customQuotaCallback = new ClientQuotaCallback {
      val quotas = new HashMap[ClientQuotaEntity, Quota]()
      override def configure(configs: Map[String, _]): Unit = {}

      override def quotaMetricTags(quotaType: ClientQuotaType, principal: KafkaPrincipal, clientId: String): Map[String, String] = Collections.emptyMap()

      override def quotaLimit(quotaType: ClientQuotaType, metricTags: Map[String, String]): java.lang.Double = 1
      override def updateClusterMetadata(cluster: Cluster): Boolean = false

      override def updateQuota(quotaType: ClientQuotaType, entity: ClientQuotaEntity, newValue: Double): Unit = {
        quotas.put(entity.asInstanceOf[KafkaQuotaEntity], new Quota(newValue.toLong, true))
      }

      override def removeQuota(quotaType: ClientQuotaType, entity: ClientQuotaEntity): Unit = {
        quotas.remove(entity.asInstanceOf[KafkaQuotaEntity])
      }

      override def quotaResetRequired(quotaType: ClientQuotaType): Boolean = false

      override def close(): Unit = {}
    }
    val clientQuotaManager = new ClientQuotaManager(
      config = new ClientQuotaManagerConfig(),
      metrics = metrics,
      quotaType = QuotaType.CONTROLLER_MUTATION,