kuwii commented on a change in pull request #35356:
URL: https://github.com/apache/spark/pull/35356#discussion_r800150649



##########
File path: 
sql/core/src/test/scala/org/apache/spark/sql/streaming/ui/StreamingQueryStatusListenerSuite.scala
##########
@@ -216,32 +216,50 @@ class StreamingQueryStatusListenerSuite extends 
StreamTest {
   }
 
   test("SPARK-38056: test writing StreamingQueryData to an in-memory store") {
-    val store = new ElementTrackingStore(new InMemoryStore(), sparkConf)
-    store.write(testStreamingQueryData)
+    testStreamingQueryData(new InMemoryStore())
   }
 
   test("SPARK-38056: test writing StreamingQueryData to a LevelDB store") {
     assume(!Utils.isMacOnAppleSilicon)
     val testDir = Utils.createTempDir()
     try {
-      val kvStore = KVUtils.open(testDir, getClass.getName)
-      val store = new ElementTrackingStore(kvStore, sparkConf)
-      store.write(testStreamingQueryData)
+      testStreamingQueryData(KVUtils.open(testDir, getClass.getName))
     } finally {
       Utils.deleteRecursively(testDir)
     }
   }
 
-  private def testStreamingQueryData: StreamingQueryData = {
-    val id = UUID.randomUUID()
-    new StreamingQueryData(
-      "some-query",
-      id,
-      id.toString,
-      isActive = false,
-      None,
-      1L,
-      None
-    )
+  test("SPARK-38056: test writing StreamingQueryData to a RocksDB store") {
+    assume(!Utils.isMacOnAppleSilicon)
+    val testDir = Utils.createTempDir()
+    try {
+      testStreamingQueryData(new RocksDB(testDir))
+    } finally {
+      Utils.deleteRecursively(testDir)
+    }
   }
-}
+
+  private def testStreamingQueryData(kvStoreFn: => KVStore): Unit = {
+    var kvStore: Option[KVStore] = None
+    try {
+      kvStore = Some(kvStoreFn)
+      val id = UUID.randomUUID()
+      val testData = new StreamingQueryData(
+        "some-query",
+        id,
+        id.toString,
+        isActive = false,
+        None,
+        1L,
+        None
+      )
+      val store = new ElementTrackingStore(kvStore.get, sparkConf)
+      try {
+        store.write(testData)
+      } finally {
+        store.close()
+      }
+    } finally {

Review comment:
       Updated in 
[a382100e21196c3362b921b559481031fbce830e](https://github.com/apache/spark/pull/35356/commits/a382100e21196c3362b921b559481031fbce830e).
 Both `ElementTrackingStore` and its `KVStore` will be both closed exactly once 
in one try-finally.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]



---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to