From b65f2d154a52275153a3fb7fdb210ba2516e2441 Mon Sep 17 00:00:00 2001
From: Dmitri Bogomolov <4glitch@gmail.com>
Date: Wed, 28 Oct 2020 18:01:55 +0200
Subject: [PATCH] Common basic cleanup procedure

---
 src/tests/common.py       | 19 +++++++++++++++++++
 src/tests/core.py         |  9 +++------
 src/tests/test_process.py |  8 +++-----
 3 files changed, 25 insertions(+), 11 deletions(-)
 create mode 100644 src/tests/common.py

diff --git a/src/tests/common.py b/src/tests/common.py
new file mode 100644
index 00000000..2f130e25
--- /dev/null
+++ b/src/tests/common.py
@@ -0,0 +1,19 @@
+import os
+
+
+_files = (
+    'keys.dat', 'debug.log', 'messages.dat', 'knownnodes.dat',
+    '.api_started', 'unittest.lock'
+)
+
+
+def cleanup(home=None, files=_files):
+    """Cleanup application files"""
+    if not home:
+        import state
+        home = state.appdata
+    for pfile in files:
+        try:
+            os.remove(os.path.join(home, pfile))
+        except OSError:
+            pass
diff --git a/src/tests/core.py b/src/tests/core.py
index 6ed46ad9..85b2cc8b 100644
--- a/src/tests/core.py
+++ b/src/tests/core.py
@@ -24,6 +24,8 @@ from network.tcp import Socks4aBMConnection, Socks5BMConnection, TCPConnection
 from queues import excQueue
 from version import softwareVersion
 
+from common import cleanup
+
 try:
     import stem.version as stem_version
 except ImportError:
@@ -48,11 +50,6 @@ def pickle_knownnodes():
         }, dst)
 
 
-def cleanup():
-    """Cleanup application files"""
-    os.remove(knownnodes_file)
-
-
 class TestCore(unittest.TestCase):
     """Test case, which runs in main pybitmessage thread"""
 
@@ -132,7 +129,7 @@ class TestCore(unittest.TestCase):
 
     def test_knownnodes_default(self):
         """test adding default knownnodes if nothing loaded"""
-        cleanup()
+        cleanup(files=('knownnodes.dat',))
         self._wipe_knownnodes()
         knownnodes.readKnownNodes()
         self.assertGreaterEqual(
diff --git a/src/tests/test_process.py b/src/tests/test_process.py
index f70c4d03..36e95b2a 100644
--- a/src/tests/test_process.py
+++ b/src/tests/test_process.py
@@ -11,6 +11,8 @@ import unittest
 
 import psutil
 
+from common import cleanup
+
 
 def put_signal_file(path, filename):
     """Creates file, presence of which is a signal about some event."""
@@ -73,11 +75,7 @@ class TestProcessProto(unittest.TestCase):
 
     @classmethod
     def _cleanup_files(cls):
-        for pfile in cls._files:
-            try:
-                os.remove(os.path.join(cls.home, pfile))
-            except OSError:
-                pass
+        cleanup(cls.home, cls._files)
 
     @classmethod
     def tearDownClass(cls):