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

Code before:
  def testQuotaTypesEnabledUpdatesOnAddAndRemove(): Unit = {
    val clientQuotaManager = new ClientQuotaManager(config, metrics, QuotaType.CONTROLLER_MUTATION, time, "")
    try {
      // Initially, quotaTypesEnabled should be QuotaTypes.NoQuotas and quotasEnabled should be false
      assertEquals(QuotaTypes.NoQuotas, clientQuotaManager.getQuotaTypesEnabled)
      assertFalse(clientQuotaManager.quotasEnabled)

      // Add a client-id quota and quotaTypesEnabled should be QuotaTypes.ClientIdQuotaEnabled
      clientQuotaManager.updateQuota(
        None,
        Some(ClientQuotaManager.ClientIdEntity("client1")),
        Some(new Quota(5, true))
      )
      assertEquals(QuotaTypes.ClientIdQuotaEnabled, clientQuotaManager.getQuotaTypesEnabled)
      assertTrue(clientQuotaManager.quotasEnabled)

      // Add a user quota and quotaTypesEnabled should be QuotaTypes.UserQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled
      clientQuotaManager.updateQuota(
        Some(ClientQuotaManager.UserEntity("userA")),
        None,
        Some(new Quota(5, true))
      )
      assertEquals(QuotaTypes.UserQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled, clientQuotaManager.getQuotaTypesEnabled)
      assertTrue(clientQuotaManager.quotasEnabled)

      // Add a user-client-id quota and quotaTypesEnabled should
      // be QuotaTypes.UserClientIdQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled | QuotaTypes.UserQuotaEnabled
      clientQuotaManager.updateQuota(
        Some(ClientQuotaManager.UserEntity("userB")),
        Some(ClientQuotaManager.ClientIdEntity("client2")),
        Some(new Quota(10, true))
      )
      assertEquals(
        QuotaTypes.UserClientIdQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled | QuotaTypes.UserQuotaEnabled,
        clientQuotaManager.getQuotaTypesEnabled
      )
      assertTrue(clientQuotaManager.quotasEnabled)

      // Remove the user quota and quotaTypesEnabled should be QuotaTypes.UserClientIdQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled
      clientQuotaManager.updateQuota(
        Some(ClientQuotaManager.UserEntity("userA")),
        None,
        None
      )
      assertEquals(QuotaTypes.UserClientIdQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled, clientQuotaManager.getQuotaTypesEnabled)
      assertTrue(clientQuotaManager.quotasEnabled)

      // Remove the client-id quota and quotaTypesEnabled should be QuotaTypes.UserClientIdQuotaEnabled
      clientQuotaManager.updateQuota(
        None,
        Some(ClientQuotaManager.ClientIdEntity("client1")),

Review comment (refactor):
could we add another case here where we add a duplicate quota type (e.g. UserClientIdQuota) and show that quotasEnabled still remains the same value?

Code after:
    }
  }

  @Test
  def testQuotaTypesEnabledUpdatesWithDefaultCallback(): Unit = {
    val clientQuotaManager = new ClientQuotaManager(config, metrics, QuotaType.CONTROLLER_MUTATION, time, "")
    try {
      // Initially, quotaTypesEnabled should be QuotaTypes.NoQuotas and quotasEnabled should be false
      assertEquals(QuotaTypes.NoQuotas, clientQuotaManager.quotaTypesEnabled)
      assertFalse(clientQuotaManager.quotasEnabled)

      // Add a client-id quota, quotaTypesEnabled should be QuotaTypes.ClientIdQuotaEnabled
      clientQuotaManager.updateQuota(None, Some(ClientQuotaManager.ClientIdEntity("client1")), Some(new Quota(5, true)))
      assertEquals(QuotaTypes.ClientIdQuotaEnabled, clientQuotaManager.quotaTypesEnabled)
      assertTrue(clientQuotaManager.quotasEnabled)

      // Add a user quota, quotaTypesEnabled should be QuotaTypes.UserQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled
      clientQuotaManager.updateQuota(Some(ClientQuotaManager.UserEntity("userA")), None, Some(new Quota(5, true)))
      assertEquals(QuotaTypes.UserQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled, clientQuotaManager.quotaTypesEnabled)
      assertTrue(clientQuotaManager.quotasEnabled)

      // Add a duplicate client-id quota, quotaTypesEnabled should remain unchanged
      clientQuotaManager.updateQuota(None, Some(ClientQuotaManager.ClientIdEntity("client2")), Some(new Quota(5, true)))
      assertEquals(QuotaTypes.UserQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled, clientQuotaManager.quotaTypesEnabled)
      assertTrue(clientQuotaManager.quotasEnabled)

      // Add duplicate user quota, quotaTypesEnabled should remain unchanged
      clientQuotaManager.updateQuota(Some(ClientQuotaManager.UserEntity("userB")), None, Some(new Quota(5, true)))
      assertEquals(QuotaTypes.UserQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled, clientQuotaManager.quotaTypesEnabled)
      assertTrue(clientQuotaManager.quotasEnabled)

      // Add a user-client-id quota, quotaTypesEnabled should be QuotaTypes.UserClientIdQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled | QuotaTypes.UserQuotaEnabled
      clientQuotaManager.updateQuota(Some(ClientQuotaManager.UserEntity("userA")), Some(ClientQuotaManager.ClientIdEntity("client1")), Some(new Quota(10, true)))
      assertEquals(QuotaTypes.UserClientIdQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled | QuotaTypes.UserQuotaEnabled, clientQuotaManager.quotaTypesEnabled)
      assertTrue(clientQuotaManager.quotasEnabled)

      // Add Duplicate user-client-id quota, quotaTypesEnabled should remain unchanged
      clientQuotaManager.updateQuota(Some(ClientQuotaManager.UserEntity("userA")), Some(ClientQuotaManager.ClientIdEntity("client1")), Some(new Quota(12, true)))
      assertEquals(QuotaTypes.UserClientIdQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled | QuotaTypes.UserQuotaEnabled, clientQuotaManager.quotaTypesEnabled)
      assertTrue(clientQuotaManager.quotasEnabled)

      // Remove the first user quota, quotaTypesEnabled should remain unchanged
      clientQuotaManager.updateQuota(Some(ClientQuotaManager.UserEntity("userA")), None, None)
      assertEquals(QuotaTypes.UserClientIdQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled | QuotaTypes.UserQuotaEnabled, clientQuotaManager.quotaTypesEnabled)
      assertTrue(clientQuotaManager.quotasEnabled)

      // Remove the second user quota, quotaTypesEnabled should be QuotaTypes.UserClientIdQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled
      clientQuotaManager.updateQuota(Some(ClientQuotaManager.UserEntity("userB")), None, None)
      assertEquals(QuotaTypes.UserClientIdQuotaEnabled | QuotaTypes.ClientIdQuotaEnabled, clientQuotaManager.quotaTypesEnabled)
      assertTrue(clientQuotaManager.quotasEnabled)