Use pytest's "mark" feature to categorise tests

Use pytest's `mark` feature to explicitly categorise the tests and run them
per category

1 files changed, 9 insertions, 10 deletions

diff --git a/setup_commands/run_tests.py b/setup_commands/run_tests.py
index 9a2c9ad..1bb5dab 100644
--- a/setup_commands/run_tests.py
+++ b/setup_commands/run_tests.py
@@ -7,21 +7,18 @@ class RunTests(Command):
     A custom command to run tests.
     """
     description = "Run the tests"
-    commands = {
-        "all": "pytest",
-        "unit": "pytest tests/unit",
-        "integration": "pytest tests/integration",
-        "performance": "pytest tests/performance",
-    }
+    test_types = (
+        "all", "unit", "integration", "performance")
     user_options = [
         ("type=", None,
          f"""Specify the type of tests to run.
-         Valid types are {tuple(commands.keys())}.
+         Valid types are {tuple(test_types)}.
          Default is `all`.""")]
 
     def __init__(self, dist):
         """Initialise the command."""
         super().__init__(dist)
+        self.command = "pytest"
 
     def initialize_options(self):
         """Initialise the default values of all the options."""
@@ -29,13 +26,15 @@ class RunTests(Command):
 
     def finalize_options(self):
         """Set final value of all the options once they are processed."""
-        if self.type not in RunTests.commands.keys():
+        if self.type not in RunTests.test_types:
             raise Exception(f"""
             Invalid test type (self.type) requested!
             Valid types are
-            {tuple(RunTests.commands.keys())}""")
+            {tuple(RunTests.test_types)}""")
 
+        if self.type != "all":
+            self.command = f"pytest -m {self.type}_test"
     def run(self):
         """Run the chosen tests"""
         print(f"Running {self.type} tests")
-        os.system(RunTests.commands[self.type])
+        os.system(self.command)