diff --git a/tests/registrator_base.py b/tests/registrator_base.py
index 309da1f..7b529e2 100644
--- a/tests/registrator_base.py
+++ b/tests/registrator_base.py
@@ -1,12 +1,11 @@
 import os
-import shutil
-import unittest
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 
 from auxiliary.turbopath import turbopath
+import shutil
 
 
-class RegistratorBase(unittest.TestCase):
+class RegistratorBase:
 
     @abstractmethod
     def get_registrator(self):
@@ -29,14 +28,12 @@ def setUp(self):
         self.moving_image = input_dir + "/tcia_example_t1.nii.gz"
 
         self.matrix = self.output_dir + f"/{self.method_name}_matrix"
-        self.transform_matrix = (
-            input_dir + f"/{self.method_name}_matrix.{self.matrix_extension}"
-        )
+        self.transform_matrix = input_dir + f"/{self.method_name}_matrix"
 
     def tearDown(self):
         # Clean up created files if they exist
-        # shutil.rmtree(self.output_dir)
-        pass
+        shutil.rmtree(self.output_dir)
+        # pass
 
     def test_register_creates_output_files(self):
         transformed_image = (
diff --git a/tests/test_registrators.py b/tests/test_registrators.py
index 6b87c60..be36fae 100644
--- a/tests/test_registrators.py
+++ b/tests/test_registrators.py
@@ -4,8 +4,10 @@
 from brainles_preprocessing.registration.eReg.eReg import eRegRegistrator
 from brainles_preprocessing.registration.niftyreg.niftyreg import NiftyRegRegistrator
 
+import unittest
 
-class TestANTsRegistrator(RegistratorBase):
+
+class TestANTsRegistrator(RegistratorBase, unittest.TestCase):
     def get_registrator(self):
         return ANTsRegistrator()
 
@@ -13,7 +15,7 @@ def get_method_and_extension(self):
         return "ants", "mat"
 
 
-class TestNiftyRegRegistratorRegistrator(RegistratorBase):
+class TestNiftyRegRegistratorRegistrator(RegistratorBase, unittest.TestCase):
     def get_registrator(self):
         return NiftyRegRegistrator()
 
@@ -21,7 +23,7 @@ def get_method_and_extension(self):
         return "niftyreg", "txt"
 
 
-class TestEregRegistratorRegistrator(RegistratorBase):
+class TestEregRegistratorRegistrator(RegistratorBase, unittest.TestCase):
     def get_registrator(self):
         return eRegRegistrator()