diff --git a/src/tests/common.py b/src/tests/common.py
index 2d60c716..4c2168b0 100644
--- a/src/tests/common.py
+++ b/src/tests/common.py
@@ -13,7 +13,10 @@ _files = (
 def cleanup(home=None, files=_files):
     """Cleanup application files"""
     if not home:
-        import state
+        try:
+            import state
+        except ImportError:
+            from pybitmessage import state
         home = state.appdata
     for pfile in files:
         try:
diff --git a/src/tests/test_protocol.py b/src/tests/test_protocol.py
index 84d87a0f..4ec4a1e7 100644
--- a/src/tests/test_protocol.py
+++ b/src/tests/test_protocol.py
@@ -3,11 +3,18 @@ Tests for common protocol functions
 """
 
 import unittest
+try:
+    from .common import cleanup
+except ImportError:
+    from common import cleanup
 
 
 class TestProtocol(unittest.TestCase):
     """Main protocol test case"""
 
+    def tearDown(self):
+        cleanup()
+
     def test_check_local(self):
         """Check the logic of TCPConnection.local"""
         from pybitmessage import protocol, state