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

Code before:
      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")),
        None
      )
      assertEquals(QuotaTypes.UserClientIdQuotaEnabled, clientQuotaManager.getQuotaTypesEnabled)
      assertTrue(clientQuotaManager.quotasEnabled)

      // Remove the user-client-id quota and quotaTypesEnabled should be QuotaTypes.NoQuotas and quotasEnabled should be false
      clientQuotaManager.updateQuota(
        Some(ClientQuotaManager.UserEntity("userB")),
        Some(ClientQuotaManager.ClientIdEntity("client2")),
        None
      )
      assertEquals(QuotaTypes.NoQuotas, clientQuotaManager.getQuotaTypesEnabled)
      assertFalse(clientQuotaManager.quotasEnabled)

Review comment (suggestion):
with the above suggestions, the remove cases become more interesting too - after removing just one `ClientIdQuota` we would expect that the quotaTypesEnabled does not change yet

Code after:
      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)

      // 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)