Add unit tests for topic emptiness class

Signed-off-by: Chase Engelbrecht <[email protected]>
engechas · Nov 6, 2023 · 3fbbcc8 · 3fbbcc8
1 parent 4fd25c7
commit 3fbbcc8
Show file tree

Hide file tree

Showing 3 changed files with 108 additions and 3 deletions.
diff --git a/.../src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java b/.../src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java
@@ -541,7 +541,7 @@ public synchronized boolean isTopicEmpty() {
         }
 
         if (currentThreadId != topicEmptinessMetadata.getTopicEmptyCheckingOwnerThreadId() ||
-                topicEmptinessMetadata.isCheckDurationExceeded(System.currentTimeMillis())) {
+                topicEmptinessMetadata.isWithinCheckInterval(System.currentTimeMillis())) {
             return topicEmptinessMetadata.isTopicEmpty();
         }
 

diff --git a/...c/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/TopicEmptinessMetadata.java b/...c/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/TopicEmptinessMetadata.java
@@ -9,7 +9,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 
 public class TopicEmptinessMetadata {
-    private static final long IS_EMPTY_CHECK_INTERVAL_MS = 60000L;
+    static final long IS_EMPTY_CHECK_INTERVAL_MS = 60000L;
 
     private long lastIsEmptyCheckTime;
     private Long topicEmptyCheckingOwnerThreadId;
@@ -41,11 +41,15 @@ public Long getTopicEmptyCheckingOwnerThreadId() {
         return this.topicEmptyCheckingOwnerThreadId;
     }
 
+    public ConcurrentHashMap<TopicPartition, Boolean> getTopicPartitionToIsEmpty() {
+        return this.topicPartitionToIsEmpty;
+    }
+
     public boolean isTopicEmpty() {
         return topicPartitionToIsEmpty.values().stream().allMatch(isEmpty -> isEmpty);
     }
 
-    public boolean isCheckDurationExceeded(final long epochTimestamp) {
+    public boolean isWithinCheckInterval(final long epochTimestamp) {
         return epochTimestamp < lastIsEmptyCheckTime + IS_EMPTY_CHECK_INTERVAL_MS;
     }
 }
diff --git a/...st/java/org/opensearch/dataprepper/plugins/kafka/consumer/TopicEmptinessMetadataTest.java b/...st/java/org/opensearch/dataprepper/plugins/kafka/consumer/TopicEmptinessMetadataTest.java
@@ -0,0 +1,101 @@
+package org.opensearch.dataprepper.plugins.kafka.consumer;
+
+import org.apache.kafka.common.TopicPartition;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+import static org.opensearch.dataprepper.plugins.kafka.consumer.TopicEmptinessMetadata.IS_EMPTY_CHECK_INTERVAL_MS;
+
+public class TopicEmptinessMetadataTest {
+    @Mock
+    private TopicPartition topicPartition;
+    @Mock
+    private TopicPartition topicPartition2;
+
+    private TopicEmptinessMetadata topicEmptinessMetadata;
+
+    @BeforeEach
+    void setup() {
+        MockitoAnnotations.openMocks(this);
+        this.topicEmptinessMetadata = new TopicEmptinessMetadata();
+    }
+
+    @Test
+    void updateTopicEmptinessStatus_AddEntry() {
+        topicEmptinessMetadata.updateTopicEmptinessStatus(topicPartition, false);
+        assertThat(topicEmptinessMetadata.getTopicPartitionToIsEmpty().containsKey(topicPartition), equalTo(true));
+        assertThat(topicEmptinessMetadata.getTopicPartitionToIsEmpty().get(topicPartition), equalTo(false));
+    }
+
+    @Test
+    void updateTopicEmptinessStatus_UpdateEntry() {
+        topicEmptinessMetadata.updateTopicEmptinessStatus(topicPartition, false);
+        assertThat(topicEmptinessMetadata.getTopicPartitionToIsEmpty().containsKey(topicPartition), equalTo(true));
+        assertThat(topicEmptinessMetadata.getTopicPartitionToIsEmpty().get(topicPartition), equalTo(false));
+
+        topicEmptinessMetadata.updateTopicEmptinessStatus(topicPartition, true);
+        assertThat(topicEmptinessMetadata.getTopicPartitionToIsEmpty().containsKey(topicPartition), equalTo(true));
+        assertThat(topicEmptinessMetadata.getTopicPartitionToIsEmpty().get(topicPartition), equalTo(true));
+    }
+
+    @Test
+    void isTopicEmpty_NoItems() {
+        assertThat(topicEmptinessMetadata.isTopicEmpty(), equalTo(true));
+    }
+
+    @Test
+    void isTopicEmpty_OnePartition_IsNotEmpty() {
+        topicEmptinessMetadata.updateTopicEmptinessStatus(topicPartition, false);
+        assertThat(topicEmptinessMetadata.isTopicEmpty(), equalTo(false));
+    }
+
+    @Test
+    void isTopicEmpty_OnePartition_IsEmpty() {
+        topicEmptinessMetadata.updateTopicEmptinessStatus(topicPartition, true);
+        assertThat(topicEmptinessMetadata.isTopicEmpty(), equalTo(true));
+    }
+
+    @Test
+    void isTopicEmpty_MultiplePartitions_OneNotEmpty() {
+        topicEmptinessMetadata.updateTopicEmptinessStatus(topicPartition, true);
+        topicEmptinessMetadata.updateTopicEmptinessStatus(topicPartition2, false);
+        assertThat(topicEmptinessMetadata.isTopicEmpty(), equalTo(false));
+    }
+
+    @Test
+    void isCheckDurationExceeded_NoPreviousChecks() {
+        assertThat(topicEmptinessMetadata.isWithinCheckInterval(System.currentTimeMillis()), equalTo(false));
+    }
+
+    @Test
+    void isCheckDurationExceeded_CurrentTimeBeforeLastCheck() {
+        final long time = System.currentTimeMillis();
+        topicEmptinessMetadata.setLastIsEmptyCheckTime(time);
+        assertThat(topicEmptinessMetadata.isWithinCheckInterval(time - 1), equalTo(true));
+    }
+
+    @Test
+    void isCheckDurationExceeded_CurrentTimeAfterLastCheck_BeforeInterval() {
+        final long time = System.currentTimeMillis();
+        topicEmptinessMetadata.setLastIsEmptyCheckTime(time);
+        assertThat(topicEmptinessMetadata.isWithinCheckInterval((time + IS_EMPTY_CHECK_INTERVAL_MS) - 1), equalTo(true));
+    }
+
+    @Test
+    void isCheckDurationExceeded_CurrentTimeAfterLastCheck_AtInterval() {
+        final long time = System.currentTimeMillis();
+        topicEmptinessMetadata.setLastIsEmptyCheckTime(time);
+        assertThat(topicEmptinessMetadata.isWithinCheckInterval(time + IS_EMPTY_CHECK_INTERVAL_MS), equalTo(false));
+    }
+
+    @Test
+    void isCheckDurationExceeded_CurrentTimeAfterLastCheck_AfterInterval() {
+        final long time = System.currentTimeMillis();
+        topicEmptinessMetadata.setLastIsEmptyCheckTime(time);
+        assertThat(topicEmptinessMetadata.isWithinCheckInterval(time + IS_EMPTY_CHECK_INTERVAL_MS + 1), equalTo(false));
+    }
+}