test/common.py

   1 """ Common housekeeping for tests that rely on daemons """
   2
   3 from configparser import ConfigParser, SectionProxy
   4 from contextlib import closing, ExitStack
   5 from http.server import HTTPServer, SimpleHTTPRequestHandler
   6 from importlib import import_module
   7 from logging import DEBUG, StreamHandler
   8 from multiprocessing import Process
   9 from os import kill, unlink
  10 from signal import SIGINT
  11 from socket import (
  12     AF_INET6,
  13     MSG_DONTWAIT,
  14     SOCK_DGRAM,
  15     SOL_SOCKET,
  16     SO_REUSEADDR,
  17     socket,
  18     SocketType,
  19 )
  20 from sys import exit, stderr
  21 from tempfile import mkstemp
  22 from time import sleep
  23 from typing import Optional
  24 from unittest import TestCase
  25
  26 NUMPORTS = 3
  27
  28
  29 class TestWithServers(TestCase):
  30     def setUp(
  31         self, *args: str, httpd: bool = False, verbose: bool = False
  32     ) -> None:
  33         freeports = []
  34         with ExitStack() as stack:
  35             for _ in range(NUMPORTS):
  36                 sk = stack.enter_context(closing(socket(AF_INET6, SOCK_DGRAM)))
  37                 sk.bind(("", 0))
  38                 sk.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
  39                 freeports.append(sk.getsockname()[1])
  40         _, self.tmpfilebase = mkstemp()
  41         self.conf = ConfigParser()
  42         self.conf["common"] = {
  43             "protocols": "zx303proto",
  44         }
  45         self.conf["collector"] = {
  46             "port": str(freeports[0]),
  47             "publishurl": "ipc://" + self.tmpfilebase + ".pub",
  48             "listenurl": "ipc://" + self.tmpfilebase + ".pul",
  49         }
  50         self.conf["storage"] = {
  51             "dbfn": self.tmpfilebase + ".storage.sqlite",
  52         }
  53         self.conf["opencellid"] = {
  54             "dbfn": self.tmpfilebase + ".opencellid.sqlite",
  55             "downloadurl": f"http://localhost:{freeports[2]}/test/262.csv.gz",
  56         }
  57         self.conf["rectifier"] = {
  58             "lookaside": "opencellid",
  59         }
  60         self.conf["wsgateway"] = {
  61             "port": str(freeports[1]),
  62         }
  63         self.children = []
  64         for srvname in args:
  65             if srvname == "collector":
  66                 kwargs = {"handle_hibernate": False}
  67             else:
  68                 kwargs = {}
  69             cls = import_module("loctrkd." + srvname, package=".")
  70             if verbose:
  71                 cls.log.addHandler(StreamHandler(stderr))
  72                 cls.log.setLevel(DEBUG)
  73             p = Process(target=cls.runserver, args=(self.conf,), kwargs=kwargs)
  74             p.start()
  75             self.children.append((srvname, p))
  76         if httpd:
  77             server = HTTPServer(("", freeports[2]), SimpleHTTPRequestHandler)
  78
  79             def run(server: HTTPServer) -> None:
  80                 try:
  81                     server.serve_forever()
  82                 except KeyboardInterrupt:
  83                     # TODO: this still leaves unclosed socket in the server
  84                     server.shutdown()
  85
  86             p = Process(target=run, args=(server,))
  87             p.start()
  88             self.children.append(("httpd", p))
  89         sleep(1)
  90
  91     def tearDown(self) -> None:
  92         for srvname, p in self.children:
  93             if p.pid is not None:
  94                 kill(p.pid, SIGINT)
  95             p.join()
  96             self.assertEqual(
  97                 p.exitcode,
  98                 0,
  99                 f"{srvname} terminated with return code {p.exitcode}",
 100             )
 101         for sfx in (
 102             "",
 103             ".pub",
 104             ".pul",
 105             ".storage.sqlite",
 106             ".opencellid.sqlite",
 107         ):
 108             try:
 109                 unlink(self.tmpfilebase + sfx)
 110             except OSError:
 111                 pass
 112
 113
 114 def send_and_drain(sock: SocketType, buf: Optional[bytes]) -> None:
 115     if buf is not None:
 116         sock.send(buf)
 117     try:
 118         sock.recv(4096, MSG_DONTWAIT)
 119     except BlockingIOError:
 120         pass