diff --git a/framework/testrun.py b/framework/testrun.py
index 4a29b4e20..df6006411 100644
--- a/framework/testrun.py
+++ b/framework/testrun.py
@@ -11,6 +11,7 @@
 import sys
 import json
 import signal
+import time
 import logger
 from device import Device
 
@@ -57,21 +58,22 @@ def start(self):
 
         self._load_devices()
 
+        self._start_network()
+
         if self._net_only:
             LOGGER.info("Network only option configured, no tests will be run")
-            self._start_network()
+            time.sleep(RUNTIME)
         else:
-            self._start_network()
-            self._start_tests()
+            self._net_orc.listener.register_callback(
+             self._device_discovered,
+             [NetworkEvent.DEVICE_DISCOVERED])
+        
+            LOGGER.info("Waiting for devices on the network...")
 
+        # Check timeout and whether testing is currently in progress before stopping
+        time.sleep(RUNTIME)
         self.stop()
 
-        # Register callbacks
-        # Disable for now as this is causing boot failures when no devices are discovered
-        # self._net_orc.listener.register_callback(
-        #     self._device_discovered,
-        #     [NetworkEvent.DEVICE_DISCOVERED])
-
     def stop(self,kill=False):
         self._stop_tests()
         self._stop_network(kill=kill)
@@ -125,9 +127,8 @@ def _get_config_abs(self,config_file=None):
     def _start_network(self):
         self._net_orc.start()
 
-    def _start_tests(self):
+    def _run_tests(self):
         """Iterate through and start all test modules."""
-
         self._test_orc.start()
 
     def _stop_network(self,kill=False):
@@ -167,6 +168,9 @@ def _device_discovered(self, mac_addr):
             LOGGER.info(
                 f'Discovered {device.make} {device.model} on the network')
         else:
+            device = Device(make=None, model=None, mac_addr=mac_addr)
             LOGGER.info(
-                f'A new device has been discovered with mac address {device.mac_addr}')
-        return device
+                f'A new device has been discovered with mac address {mac_addr}')
+        
+        # TODO: Pass device information to test orchestrator/runner
+        self._run_tests()
diff --git a/test_orc/modules/baseline/bin/start_test_module b/test_orc/modules/baseline/bin/start_test_module
index 292b57de2..2938eb0f8 100644
--- a/test_orc/modules/baseline/bin/start_test_module
+++ b/test_orc/modules/baseline/bin/start_test_module
@@ -37,4 +37,6 @@ chown $HOST_USER:$HOST_USER $RESULT_FILE
 # Run the python scrip that will execute the tests for this module
 # -u flag allows python print statements
 # to be logged by docker by running unbuffered
-python3 -u $PYTHON_SRC_DIR/run.py "-m $MODULE_NAME"
\ No newline at end of file
+python3 -u $PYTHON_SRC_DIR/run.py "-m $MODULE_NAME"
+
+echo Module has finished
\ No newline at end of file
diff --git a/test_orc/modules/baseline/python/src/test_module.py b/test_orc/modules/baseline/python/src/test_module.py
index 440b87f7f..d4065cde3 100644
--- a/test_orc/modules/baseline/python/src/test_module.py
+++ b/test_orc/modules/baseline/python/src/test_module.py
@@ -32,8 +32,6 @@ def run_tests(self):
         self.module_test2 = False
         LOGGER.info("Test 2 complete.")
 
-        time.sleep(10)
-
     def generate_results(self):
         results = []
         results.append(self.generate_result("Test 1", self.module_test1))
diff --git a/test_orc/python/src/module.py b/test_orc/python/src/module.py
index 6d24d7e1e..8121c34db 100644
--- a/test_orc/python/src/module.py
+++ b/test_orc/python/src/module.py
@@ -1,6 +1,6 @@
 """Represemts a test module."""
 from dataclasses import dataclass
-from docker.client.Container import Container
+from docker.models.containers import Container
 
 @dataclass
 class TestModule: # pylint: disable=too-few-public-methods,too-many-instance-attributes
diff --git a/test_orc/python/src/test_orchestrator.py b/test_orc/python/src/test_orchestrator.py
index 77f73f407..f68a13579 100644
--- a/test_orc/python/src/test_orchestrator.py
+++ b/test_orc/python/src/test_orchestrator.py
@@ -6,9 +6,10 @@
 import docker
 from docker.types import Mount
 import logger
+from module import TestModule
 
 LOG_NAME = "test_orc"
-LOGGER = logger.get_logger('test_orc')
+LOGGER = logger.get_logger("test_orc")
 RUNTIME_DIR = "runtime"
 TEST_MODULES_DIR = "modules"
 MODULE_CONFIG = "conf/module_config.json"
@@ -196,24 +197,4 @@ def _stop_module(self, module, kill=False):
                     container.stop()
                 LOGGER.debug("Container stopped:" + module.container_name)
         except docker.errors.NotFound:
-            pass
-
-class TestModule: # pylint: disable=too-few-public-methods,too-many-instance-attributes
-    """Represents a test module."""
-
-    def __init__(self):
-        self.name = None
-        self.display_name = None
-        self.description = None
-
-        self.build_file = None
-        self.container = None
-        self.container_name = None
-        self.image_name = None
-        self.enable_container = True
-
-        self.timeout = 60
-
-        # Absolute path
-        self.dir = None
-        self.dir_name = None
+            pass
\ No newline at end of file