[Test] Small refactoring in dummy services test

Moved code for pid writing and terminating into a separated module.
Also added cleanup: in some cases, processes remained in the system after test is done. It should not happen anymore

6 files changed, 78 insertions, 30 deletions

diff --git a/test/functional/cases/160_antivirus.robot b/test/functional/cases/160_antivirus.robot
index 1516ae1b5..86c2e51b0 100644
--- a/test/functional/cases/160_antivirus.robot
+++ b/test/functional/cases/160_antivirus.robot
@@ -19,12 +19,14 @@ CLAMAV MISS
   Run Dummy Clam  ${PORT_CLAM}
   ${result} =  Scan Message With Rspamc  ${MESSAGE}
   Check Rspamc  ${result}  CLAM_VIRUS  inverse=1
+  Shutdown clamav
 
 CLAMAV HIT
   Run Dummy Clam  ${PORT_CLAM}  1
   ${result} =  Scan Message With Rspamc  ${MESSAGE2}
   Check Rspamc  ${result}  CLAM_VIRUS (1.00)[Eicar-Test-Signature]
   Should Not Contain  ${result.stdout}  CLAMAV_FAIL
+  Shutdown clamav
 
 CLAMAV CACHE HIT
   ${result} =  Scan Message With Rspamc  ${MESSAGE2}
@@ -41,6 +43,7 @@ FPROT MISS
   ${result} =  Scan Message With Rspamc  ${MESSAGE2}
   Check Rspamc  ${result}  FPROT_VIRUS  inverse=1
   Should Not Contain  ${result.stdout}  FPROT_EICAR
+  Shutdown fport
 
 FPROT HIT - PATTERN
   Run Dummy Fprot  ${PORT_FPROT}  1
@@ -52,6 +55,8 @@ FPROT HIT - PATTERN
   Should Contain  ${result.stdout}  FPROT_VIRUS_DUPLICATE_PATTERN
   Should Not Contain  ${result.stdout}  FPROT_VIRUS_DUPLICATE_DEFAULT
   Should Not Contain  ${result.stdout}  FPROT_VIRUS_DUPLICATE_NOPE
+  Shutdown fport
+  Shutdown fport duplicate
 
 FPROT CACHE HIT
   ${result} =  Scan Message With Rspamc  ${MESSAGE}
@@ -75,6 +80,21 @@ Antivirus Setup
 Antivirus Teardown
   Normal Teardown
   Shutdown Process With Children  ${REDIS_PID}
+  Shutdown clamav
+  Shutdown fport
+  Terminate All Processes    kill=True
+
+Shutdown clamav
+  ${clamav_pid} =  Get File if exists  /tmp/dummy_clamav.pid
+  Run Keyword if  ${clamav_pid}  Shutdown Process With Children  ${clamav_pid}
+
+Shutdown fport
+  ${fport_pid} =  Get File if exists  /tmp/dummy_fprot.pid
+  Run Keyword if  ${fport_pid}  Shutdown Process With Children  ${fport_pid}
+
+Shutdown fport duplicate
+  ${fport_pid} =  Get File if exists  /tmp/dummy_fprot_dupe.pid
+  Run Keyword if  ${fport_pid}  Shutdown Process With Children  ${fport_pid}
 
 Run Dummy Clam
   [Arguments]  ${port}  ${found}=
@@ -83,5 +103,5 @@ Run Dummy Clam
 
 Run Dummy Fprot
   [Arguments]  ${port}  ${found}=  ${pid}=/tmp/dummy_fprot.pid
-  ${result} =  Start Process  ${TESTDIR}/util/dummy_fprot.py  ${port}  ${found}  ${pid}
+  Start Process  ${TESTDIR}/util/dummy_fprot.py  ${port}  ${found}  ${pid}
   Wait Until Created  ${pid}
diff --git a/test/functional/lib/rspamd.py b/test/functional/lib/rspamd.py
index 97ab516c1..5673a10d9 100644
--- a/test/functional/lib/rspamd.py
+++ b/test/functional/lib/rspamd.py
@@ -212,3 +212,8 @@ def shutdown_process_with_children(pid):
         except:
             pass
 
+def get_file_if_exists(file_path):
+    if os.path.exists(file_path):
+        with open(file_path, 'r') as myfile:
+            return myfile.read()
+    return None
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)