Repository: apache/kafka
File: server/src/test/java/org/apache/kafka/server/quota/ThrottledChannelExpirationTest.java (Java)

Code before:
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.kafka.server.quota;

import org.apache.kafka.common.metrics.MetricConfig;
import org.apache.kafka.common.metrics.Metrics;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.server.config.ClientQuotaManagerConfig;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.util.Collections;
import java.util.concurrent.DelayQueue;
import java.util.concurrent.TimeUnit;

import static org.junit.jupiter.api.Assertions.assertEquals;

public class ThrottledChannelExpirationTest {
    private MockTime time;
    private int numCallbacksForStartThrottling;
    private int numCallbacksForEndThrottling;
    private Metrics metrics;
    private ThrottleCallback callback;

    @BeforeEach
    public void beforeMethod() {
        time = new MockTime();
        numCallbacksForStartThrottling = 0;
        numCallbacksForEndThrottling = 0;
        metrics = new Metrics(new MetricConfig(), Collections.emptyList(), time);
        callback = new ThrottleCallback() {
            @Override
            public void startThrottling() {
                numCallbacksForStartThrottling++;
            }

            @Override
            public void endThrottling() {
                numCallbacksForEndThrottling++;
            }
        };
    }

    @Test

Review comment (suggestion):
Would you mind initializing these in the constructor?
```java
    private final MockTime time = new MockTime();
    private int numCallbacksForStartThrottling = 0;
    private int numCallbacksForEndThrottling = 0;
    private final Metrics metrics = new Metrics(new MetricConfig(), List.of(), time);
    private final ThrottleCallback callback = new ThrottleCallback() {
        @Override
        public void startThrottling() {
            numCallbacksForStartThrottling++;
        }

        @Override
        public void endThrottling() {
            numCallbacksForEndThrottling++;
        }
    };
```

Code after:
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.kafka.server.quota;

import org.apache.kafka.common.metrics.MetricConfig;
import org.apache.kafka.common.metrics.Metrics;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.server.config.ClientQuotaManagerConfig;

import org.junit.jupiter.api.Test;

import java.util.List;
import java.util.concurrent.DelayQueue;
import java.util.concurrent.TimeUnit;

import static org.junit.jupiter.api.Assertions.assertEquals;

public class ThrottledChannelExpirationTest {
    private final MockTime time = new MockTime();
    private final Metrics metrics = new Metrics(new MetricConfig(), List.of(), time);
    private int numCallbacksForStartThrottling = 0;
    private int numCallbacksForEndThrottling = 0;
    private final ThrottleCallback callback = new ThrottleCallback() {
        @Override
        public void startThrottling() {
            numCallbacksForStartThrottling++;
        }

        @Override
        public void endThrottling() {
            numCallbacksForEndThrottling++;
        }
    };

    @Test
    public void testCallbackInvocationAfterExpiration() {
        ClientQuotaManager clientMetrics = new ClientQuotaManager(new ClientQuotaManagerConfig(), metrics, QuotaType.PRODUCE, time, "");

        DelayQueue<ThrottledChannel> delayQueue = new DelayQueue<>();
        var reaper = clientMetrics.new ThrottledChannelReaper(delayQueue, "");
        try {
            // Add 4 elements to the queue out of order. Add 2 elements with the same expire timestamp.
            ThrottledChannel channel1 = new ThrottledChannel(time, 10, callback);
            ThrottledChannel channel2 = new ThrottledChannel(time, 30, callback);
            ThrottledChannel channel3 = new ThrottledChannel(time, 30, callback);