4 files changed, 52 insertions, 29 deletions
diff --git a/test/functional/util/dummy_clam.py b/test/functional/util/dummy_clam.py
index 6be20de57..d25a63525 100755
--- a/test/functional/util/dummy_clam.py
+++ b/test/functional/util/dummy_clam.py
@@ -4,16 +4,16 @@ PID = "/tmp/dummy_clamav.pid"
 
 import os
 import sys
+import socket
+import dummy_killer
 try:
     import SocketServer as socketserver
 except:
     import socketserver
-import signal
 
 class MyTCPHandler(socketserver.BaseRequestHandler):
 
     def handle(self):
-        os.remove(PID)
         self.data = self.request.recv(1024).strip()
         if self.server.foundvirus:
             self.request.sendall(b"stream: Eicar-Test-Signature FOUND\0")
@@ -22,11 +22,6 @@ class MyTCPHandler(socketserver.BaseRequestHandler):
         self.request.close()
 
 if __name__ == "__main__":
-    pid = os.fork()
-    if pid > 0:
-        sys.exit(0)
-    signal.alarm(5)
-
     HOST = "localhost"
 
     alen = len(sys.argv)
@@ -45,7 +40,13 @@ if __name__ == "__main__":
     server.foundvirus = foundvirus
     server.server_bind()
     server.server_activate()
-    open(PID, 'w').close()
-    server.handle_request()
+
+    dummy_killer.setup_killer(server)
+    dummy_killer.write_pid(PID)
+
+    try:
+        server.handle_request()
+    except socket.error:
+        print "Socket closed"
+
     server.server_close()
-    os.exit(0)
diff --git a/test/functional/util/dummy_fprot.py b/test/functional/util/dummy_fprot.py
index 34725280b..e6dc58bdd 100755
--- a/test/functional/util/dummy_fprot.py
+++ b/test/functional/util/dummy_fprot.py
@@ -2,6 +2,8 @@
 import os
 import sys
 import signal
+import socket
+import dummy_killer
 
 
 try:
@@ -14,7 +16,6 @@ PID = "/tmp/dummy_fprot.pid"
 class MyTCPHandler(socketserver.BaseRequestHandler):
 
     def handle(self):
-        os.remove(PID)
         self.data = self.request.recv(1024).strip()
         if self.server.foundvirus:
             self.request.sendall(b"1 <infected: EICAR_Test_File> FOO->bar\n")
@@ -23,10 +24,6 @@ class MyTCPHandler(socketserver.BaseRequestHandler):
         self.request.close()
 
 if __name__ == "__main__":
-    pid = os.fork()
-    if pid > 0:
-        sys.exit(0)
-    signal.alarm(5)
 
     HOST = "localhost"
 
@@ -49,7 +46,12 @@ if __name__ == "__main__":
     server.foundvirus = foundvirus
     server.server_bind()
     server.server_activate()
-    open(PID, 'w').close()
-    server.handle_request()
+
+    dummy_killer.setup_killer(server)
+    dummy_killer.write_pid(PID)
+
+    try:
+        server.handle_request()
+    except socket.error:
+        print "Socket closed"
     server.server_close()
-    os.exit(0)
diff --git a/test/functional/util/dummy_http.py b/test/functional/util/dummy_http.py
index 4814613ea..dc4cee354 100755
--- a/test/functional/util/dummy_http.py
+++ b/test/functional/util/dummy_http.py
@@ -3,11 +3,11 @@
 import BaseHTTPServer
 import SocketServer
 import SimpleHTTPServer
+import dummy_killer
 
 import time
 import os
 import sys
-import signal
 import socket
 
 PORT = 18080
@@ -90,9 +90,7 @@ class ThreadingSimpleServer(SocketServer.ThreadingMixIn,
         self.timeout = 1
         
     def run(self):
-        with open(PID, 'w+') as f:
-            f.write(str(os.getpid()))
-            f.close()
+        dummy_killer.write_pid(PID)
         try:
             while 1:
                 sys.stdout.flush()
@@ -101,7 +99,6 @@ class ThreadingSimpleServer(SocketServer.ThreadingMixIn,
             print "Interrupt"
         except socket.error:
             print "Socket closed"
-            pass
 
     def stop(self):
         self.keep_running = False
@@ -111,11 +108,6 @@ class ThreadingSimpleServer(SocketServer.ThreadingMixIn,
 if __name__ == '__main__':
     server = ThreadingSimpleServer()
 
-    def alarm_handler(signum, frame):
-        server.stop()
-
-    signal.signal(signal.SIGALRM, alarm_handler)
-    signal.signal(signal.SIGTERM, alarm_handler)
-    signal.alarm(1000)
+    dummy_killer.setup_killer(server, server.stop)
 
     server.run()
diff --git a/test/functional/util/dummy_killer.py b/test/functional/util/dummy_killer.py
new file mode 100644
index 000000000..723e6ef4a
--- /dev/null
+++ b/test/functional/util/dummy_killer.py
@@ -0,0 +1,28 @@
+import signal
+import os
+import atexit
+
+def setup_killer(server, method = None):
+    def default_method():
+        server.server_close()
+    
+    if method is None:
+        method = default_method
+
+    def alarm_handler(signum, frame):
+        method()
+
+    signal.signal(signal.SIGALRM, alarm_handler)
+    signal.signal(signal.SIGTERM, alarm_handler)
+    signal.alarm(10)
+
+
+def write_pid(path):
+    with open(path, 'w+') as f:
+        f.write(str(os.getpid()))
+        f.close()
+
+    def cleanup():
+        os.remove(path)
+
+    atexit.register(cleanup)