diff --git a/test_daemon.py b/test_daemon.py
index f0e486b..ad1ef1d 100755
--- a/test_daemon.py
+++ b/test_daemon.py
@@ -6,49 +6,78 @@
 import fcntl
 import subprocess
 
-def main():
-    socket_path = "./test.sock"
-    daemon_command = "./.lake/build/bin/sand"
-    daemon_args = ["daemon"]
+from contextlib import contextmanager
 
-    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-    
-    # Remove the socket file if it already exists
+SOCKET_PATH = "./test.sock"
+
+'''
+Remove the socket file if it already exists
+'''
+def ensure_socket_deleted():
     try:
-        os.unlink(socket_path)
-    except OSError:
-        if os.path.exists(socket_path):
-            raise
+        os.unlink(SOCKET_PATH)
+    except FileNotFoundError:
+        pass
+
+'''
+Ensures when the context manager exits:
+- the daemon process is terminated 
+- the socket file is removed
 
-    sock.bind(socket_path)
+(assuming we're not killed with SIGKILL)
+'''
+@contextmanager
+def daemon():
+    ensure_socket_deleted()
+
+    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    sock.bind(SOCKET_PATH)
     sock.listen(1)
 
     flags = fcntl.fcntl(sock.fileno(), fcntl.F_GETFD)
     flags |= fcntl.FD_CLOEXEC
     fcntl.fcntl(sock.fileno(), fcntl.F_SETFD, flags)
 
-    print(f"Python: Socket created at {socket_path} on fd {sock.fileno()}")
+    print(f"-- Socket created at {SOCKET_PATH} on fd {sock.fileno()}")
 
-    daemon_proc = subprocess.Popen(
-        [daemon_command] + daemon_args,
-        pass_fds=(sock.fileno(),),
-    )
+    daemon_command = "./.lake/build/bin/sand"
+    daemon_args = ["daemon"]
+    try:
+        daemon_proc = subprocess.Popen(
+            [daemon_command] + daemon_args,
+            pass_fds=(sock.fileno(),),
+        )
 
-    print(f"Python: Daemon started with PID {daemon_proc.pid}")
+        print(f"-- Daemon started with PID {daemon_proc.pid}")
+        # Close the socket in the parent process
+        sock.close()
+        yield daemon_proc
+    finally:
+        print(f"-- Terminating daemon with PID {daemon_proc.pid}")
+        daemon_proc.terminate()
+        print(f"-- Waiting for daemon to terminate")
+        daemon_proc.wait()
+        print(f"-- Daemon terminated")
+        print(f"-- Removing socket file {SOCKET_PATH}")
+        ensure_socket_deleted()
+        
 
-    # Close the socket in the parent process
-    sock.close()
+def main():
+    print("--------------------------")
+    print("Starting integration tests")
+    print("--------------------------")
 
-    run_client_tests(socket_path)
+    with daemon() as daemon_proc:
+        run_client_tests()
 
-def run_client_tests(socket_path):
-    print(f"Python: Running client tests against {socket_path}")
+def run_client_tests():
+    print(f"-- Running client tests against {SOCKET_PATH}")
 
     client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-    client_sock.connect(socket_path)
+    client_sock.connect(SOCKET_PATH)
     client_sock.send(b'"list"')
     response = client_sock.recv(1024)
-    print(f"Python: Received response: {response}")
+    print(f"-- Received response: {response}")
     client_sock.close()
 
 if __name__ == "__main__":