diff --git a/node/src/integration-test/kotlin/net/corda/node/NodeUnloadHandlerTests.kt b/node/src/test/kotlin/net/corda/node/internal/NodeUnloadHandlerTests.kt
similarity index 55%
rename from node/src/integration-test/kotlin/net/corda/node/NodeUnloadHandlerTests.kt
rename to node/src/test/kotlin/net/corda/node/internal/NodeUnloadHandlerTests.kt
index 0efbcf12a1..295c6ede59 100644
--- a/node/src/integration-test/kotlin/net/corda/node/NodeUnloadHandlerTests.kt
+++ b/node/src/test/kotlin/net/corda/node/internal/NodeUnloadHandlerTests.kt
@@ -1,13 +1,12 @@
-package net.corda.node
+package net.corda.node.internal
 
+import net.corda.core.internal.packageName
 import net.corda.core.node.ServiceHub
 import net.corda.core.node.services.CordaService
 import net.corda.core.serialization.SingletonSerializeAsToken
 import net.corda.core.utilities.contextLogger
-import net.corda.core.utilities.getOrThrow
-import net.corda.testing.core.DUMMY_BANK_A_NAME
-import net.corda.testing.driver.DriverParameters
-import net.corda.testing.driver.driver
+import net.corda.testing.node.internal.InternalMockNetwork
+import org.junit.After
 import org.junit.Assert.assertTrue
 import org.junit.Test
 import java.util.concurrent.CountDownLatch
@@ -15,16 +14,23 @@ import java.util.concurrent.TimeUnit
 
 class NodeUnloadHandlerTests {
     companion object {
-        val latch = CountDownLatch(1)
+        val registerLatch = CountDownLatch(1)
+        val shutdownLatch = CountDownLatch(1)
+    }
+
+    private val mockNet = InternalMockNetwork(cordappPackages = listOf(javaClass.packageName), notarySpecs = emptyList())
+
+    @After
+    fun cleanUp() {
+        mockNet.stopNodes()
     }
 
     @Test
     fun `should be able to register run on stop lambda`() {
-        driver(DriverParameters(startNodesInProcess = true, extraCordappPackagesToScan = listOf("net.corda.node"), notarySpecs = emptyList())) {
-            startNode(providedName = DUMMY_BANK_A_NAME).getOrThrow()
-            // just want to fall off the end of this for the mo...
-        }
-        assertTrue("Timed out waiting for AbstractNode to invoke the test service shutdown callback", latch.await(30, TimeUnit.SECONDS))
+        val node = mockNet.createNode()
+        registerLatch.await()  // Make sure the handler is registered on node start up
+        node.dispose()
+        assertTrue("Timed out waiting for AbstractNode to invoke the test service shutdown callback", shutdownLatch.await(30, TimeUnit.SECONDS))
     }
 
     @Suppress("unused")
@@ -36,11 +42,12 @@ class NodeUnloadHandlerTests {
 
         init {
             serviceHub.registerUnloadHandler(this::shutdown)
+            registerLatch.countDown()
         }
 
         private fun shutdown() {
             log.info("shutting down")
-            latch.countDown()
+            shutdownLatch.countDown()
         }
     }
 }