This repository has been archived on 2023-03-26. You can view files and clone it, but cannot push or open issues or pull requests.
catfish/tests/__init__.py

85 lines
2.5 KiB
Python
Raw Normal View History

2018-12-12 22:34:51 +00:00
import functools
import os
import shutil
import subprocess
import time
2018-12-18 15:49:20 +00:00
from pathlib import Path
from unittest import TestCase
2018-12-12 22:34:51 +00:00
import psutil
2018-12-15 16:08:37 +00:00
from aiohttp.test_utils import AioHTTPTestCase, unused_port
2018-12-12 22:34:51 +00:00
from aiounittest import AsyncTestCase
2018-12-12 15:53:56 +00:00
from click.testing import CliRunner
2018-12-12 22:34:51 +00:00
from catfish.__main__ import cli
2018-12-15 16:08:37 +00:00
from catfish.router import get_server
from catfish.utils.processes import (
get_root_process,
terminate_processes,
terminate_subprocesses,
)
2018-12-13 11:44:41 +00:00
from catfish.utils.sockets import create_base_socket_dir, delete_base_socket_dir
from catfish.worker.server import WORKER_SERVER_SOCKET
2018-12-12 15:53:56 +00:00
class BaseTestCase(TestCase):
2018-12-18 15:49:20 +00:00
TESTS_DIR = Path(os.path.dirname(__file__))
EXAMPLE_DIR = TESTS_DIR.parent.joinpath("example")
2018-12-18 20:16:34 +00:00
DUMMY_EXE = EXAMPLE_DIR.joinpath("src", "dummy_program.py")
2018-12-12 15:53:56 +00:00
def setUp(self):
2018-12-13 11:44:41 +00:00
create_base_socket_dir()
2018-12-12 15:53:56 +00:00
self.cli_runner = CliRunner()
self.cli = cli
2018-12-12 18:47:40 +00:00
self.run_cli = functools.partial(self.cli_runner.invoke, self.cli)
def terminate_dummy_processes(self):
dummy_processes = []
for process in get_root_process().children(recursive=True):
try:
if str(self.DUMMY_EXE) in " ".join(process.cmdline()):
dummy_processes.append(process)
except psutil.NoSuchProcess:
continue
terminate_processes(dummy_processes)
2018-12-12 18:47:40 +00:00
def tearDown(self):
terminate_subprocesses()
self.terminate_dummy_processes()
2018-12-13 11:44:41 +00:00
delete_base_socket_dir()
2018-12-12 21:55:45 +00:00
class BaseWorkerTestCase(BaseTestCase):
def setUp(self):
super().setUp()
self.unused_port = unused_port()
self.worker_process = subprocess.Popen(
[
shutil.which("ctf"),
"start",
"--port",
str(self.unused_port),
"--no-fork",
],
2018-12-14 22:43:20 +00:00
universal_newlines=True,
)
while not WORKER_SERVER_SOCKET.exists():
time.sleep(0.1)
2018-12-12 21:55:45 +00:00
def tearDown(self):
self.worker_process.terminate()
self.worker_process.wait()
2018-12-12 21:55:45 +00:00
super().tearDown()
2018-12-15 16:08:37 +00:00
class BaseRouterTestCase(AsyncTestCase, AioHTTPTestCase):
async def get_application(self):
return get_server()
def __getattribute__(self, name):
attr = super(AioHTTPTestCase, self).__getattribute__(name)
if name.startswith("test_") and callable(attr):
return lambda: self.loop.run_until_complete(attr())
else:
return attr