From ad637fbbdc3acc9e8abda14a1a2b560d55d5ca5a Mon Sep 17 00:00:00 2001
From: Demogorgon314 <kwang@streamnative.io>
Date: Mon, 6 Dec 2021 12:03:58 +0800
Subject: [PATCH] Use MetadataStore instead of mockZooKeeper in
 ProducerIdManagerTest

---
 .../transaction/ProducerIdManagerTest.java    | 28 ++++++++++---------
 1 file changed, 15 insertions(+), 13 deletions(-)

diff --git a/tests/src/test/java/io/streamnative/pulsar/handlers/kop/coordinator/transaction/ProducerIdManagerTest.java b/tests/src/test/java/io/streamnative/pulsar/handlers/kop/coordinator/transaction/ProducerIdManagerTest.java
index d165028955..0ca353495a 100644
--- a/tests/src/test/java/io/streamnative/pulsar/handlers/kop/coordinator/transaction/ProducerIdManagerTest.java
+++ b/tests/src/test/java/io/streamnative/pulsar/handlers/kop/coordinator/transaction/ProducerIdManagerTest.java
@@ -17,8 +17,9 @@
 import static org.testng.AssertJUnit.fail;
 
 import io.streamnative.pulsar.handlers.kop.KopProtocolHandlerTestBase;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
 import lombok.extern.slf4j.Slf4j;
-import org.apache.zookeeper.data.Stat;
 import org.testng.Assert;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
@@ -31,6 +32,8 @@
 @Slf4j
 public class ProducerIdManagerTest extends KopProtocolHandlerTestBase {
 
+    private static final long DEFAULT_TEST_TIMEOUT = 20 * 1000;
+
     @BeforeClass
     @Override
     protected void setup() throws Exception {
@@ -43,15 +46,13 @@ protected void cleanup() throws Exception {
         super.internalCleanup();
     }
 
-    @BeforeMethod
+    @BeforeMethod(timeOut = DEFAULT_TEST_TIMEOUT)
     protected void cleanZNode() throws Exception {
-        Stat stat = mockZooKeeper.exists(ProducerIdManager.KOP_PID_BLOCK_ZNODE, null);
-        if (stat != null) {
-            mockZooKeeper.delete(ProducerIdManager.KOP_PID_BLOCK_ZNODE, -1);
-        }
+        pulsar.getLocalMetadataStore()
+                .deleteRecursive(ProducerIdManager.KOP_PID_BLOCK_ZNODE).get(10, TimeUnit.SECONDS);
     }
 
-    @Test
+    @Test(timeOut = DEFAULT_TEST_TIMEOUT)
     public void testGetProducerId() throws Exception {
         ProducerIdManager manager1 = new ProducerIdManager(0, pulsar.getLocalMetadataStore());
         manager1.initialize().get();
@@ -76,13 +77,14 @@ public void testGetProducerId() throws Exception {
         assertEquals(pid2 + ProducerIdManager.PID_BLOCK_SIZE * 2, manager2.generateProducerId().get().longValue());
     }
 
-    @Test
+    @Test(timeOut = DEFAULT_TEST_TIMEOUT)
     public void testExceedProducerIdLimit() throws Exception {
-        mockZooKeeper.create(ProducerIdManager.KOP_PID_BLOCK_ZNODE, null, null, null);
-        mockZooKeeper.setData(ProducerIdManager.KOP_PID_BLOCK_ZNODE,
-                ProducerIdManager.generateProducerIdBlockJson(
-                        new ProducerIdManager.ProducerIdBlock(
-                                1, Long.MAX_VALUE - ProducerIdManager.PID_BLOCK_SIZE, Long.MAX_VALUE)), -1);
+        pulsar.getLocalMetadataStore()
+                .put(ProducerIdManager.KOP_PID_BLOCK_ZNODE,
+                        ProducerIdManager.generateProducerIdBlockJson(
+                                new ProducerIdManager.ProducerIdBlock(
+                                        1, Long.MAX_VALUE - ProducerIdManager.PID_BLOCK_SIZE, Long.MAX_VALUE)),
+                        Optional.empty()).get(10, TimeUnit.SECONDS);
 
         ProducerIdManager producerIdManager = new ProducerIdManager(0, pulsar.getLocalMetadataStore());
         try {