diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/store/rocksdb/RocksDBStorageEngineTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/store/rocksdb/RocksDBStorageEngineTest.java index eafdae7b50..ffa2bfd406 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/store/rocksdb/RocksDBStorageEngineTest.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/store/rocksdb/RocksDBStorageEngineTest.java @@ -22,6 +22,7 @@ import java.util.Set; import org.testng.Assert; import org.testng.annotations.AfterClass; +import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; @@ -34,6 +35,7 @@ public class RocksDBStorageEngineTest extends AbstractStorageEngineTest { private final ReadOnlyStoreRepository mockReadOnlyStoreRepository = mock(ReadOnlyStoreRepository.class); private static final int versionNumber = 0; private static final String topicName = Version.composeKafkaTopic(storeName, versionNumber); + private int testCount = 0; @Override public void createStorageEngineForTest() { @@ -67,6 +69,24 @@ public void cleanUp() throws Exception { storageService.stop(); } + @AfterMethod + public void testCounter() { + this.testCount++; + } + + /** + * Some tests require a reset if other tests have run before them, as they are sensitive to contamination. + * + * Alternatively, we could make {@link #setUp()} have {@link org.testng.annotations.BeforeMethod} and + * {@link #cleanUp()} have {@link AfterMethod}, though that makes the class take longer than the current approach. + */ + private void reset() throws Exception { + if (this.testCount > 0) { + cleanUp(); + setUp(); + } + } + @Test public void testGetAndPut() { super.testGetAndPut(); @@ -154,11 +174,13 @@ public void testGetInvalidKeys() { @Test public void testPartitioning() throws Exception { + reset(); super.testPartitioning(); } @Test public void testAddingAPartitionTwice() throws Exception { + reset(); super.testAddingAPartitionTwice(); }