From 338c006bb15e75c1f62cc1b9b7aa8f295d883b2f Mon Sep 17 00:00:00 2001
From: Lee Miller <lee.miller@tutanota.com>
Date: Thu, 7 Dec 2023 00:27:04 +0200
Subject: [PATCH] Inherit TestProofofworkBase from partial.TestPartialRun

---
 src/tests/partial.py          |  3 +++
 src/tests/test_proofofwork.py | 19 +++++++------------
 2 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/src/tests/partial.py b/src/tests/partial.py
index 870f6626..f97dc414 100644
--- a/src/tests/partial.py
+++ b/src/tests/partial.py
@@ -2,6 +2,7 @@
 
 import os
 import sys
+import time
 import unittest
 
 from pybitmessage import pathmagic
@@ -39,3 +40,5 @@ class TestPartialRun(unittest.TestCase):
         # deactivate pathmagic
         os.chdir(cls.dirs[0])
         sys.path.remove(cls.dirs[1])
+        time.sleep(5)
+        cls.state.shutdown = 0
diff --git a/src/tests/test_proofofwork.py b/src/tests/test_proofofwork.py
index d3c1d6fd..2bf4fa08 100644
--- a/src/tests/test_proofofwork.py
+++ b/src/tests/test_proofofwork.py
@@ -14,17 +14,22 @@ from pybitmessage.defaults import (
     networkDefaultProofOfWorkNonceTrialsPerByte,
     networkDefaultPayloadLengthExtraBytes)
 
+from .partial import TestPartialRun
 from .samples import sample_pow_target, sample_pow_initial_hash
 
 default_ttl = 7200
 
 
-class TestProofofworkBase(unittest.TestCase):
+class TestProofofworkBase(TestPartialRun):
     """Basic test case for proofofwork"""
 
     @classmethod
     def setUpClass(cls):
         proofofwork.init()
+        super(TestProofofworkBase, cls).setUpClass()
+
+    def setUp(self):
+        self.state.shutdown = 0
 
     @staticmethod
     def _make_sample_payload(TTL=default_ttl):
@@ -43,14 +48,6 @@ class TestProofofworkBase(unittest.TestCase):
 class TestProofofwork(TestProofofworkBase):
     """The main test case for proofofwork"""
 
-    @classmethod
-    def tearDownClass(cls):
-        import state
-        state.shutdown = 0
-
-    def setUp(self):
-        self.tearDownClass()
-
     def _make_sample_data(self):
         payload = self._make_sample_payload()
         return payload, proofofwork.getTarget(
@@ -70,10 +67,8 @@ class TestProofofwork(TestProofofworkBase):
                 pack('>Q', nonce) + payload, 2000, 2000,
                 int(time.time()) + TTL - 3600))
 
-        import state
-
         with self.assertRaises(StopIteration):
-            state.shutdown = 1
+            self.state.shutdown = 1
             proofofwork.calculate(payload, TTL)
 
     def test_CPoW(self):