catfish/tests/__init__.py

import functools
import os
import shutil
import subprocess

from aiohttp.test_utils import unused_port
from aiounittest import AsyncTestCase
from click.testing import CliRunner

from catfish.__main__ import cli
from catfish.utils.processes import CURRENT_PROCESS, terminate_processes
from catfish.utils.sockets import create_base_socket_dir, delete_base_socket_dir


class BaseTestCase(AsyncTestCase):
    DUMMY_EXE = os.path.join(os.path.dirname(__file__), "dummy_program.py")

    def setUp(self):
        create_base_socket_dir()
        self.cli_runner = CliRunner()
        self.cli = cli
        self.run_cli = functools.partial(self.cli_runner.invoke, self.cli)

    def tearDown(self):
        terminate_processes(CURRENT_PROCESS.children(recursive=True))
        delete_base_socket_dir()


class BaseWorkerTestCase(BaseTestCase):
    def setUp(self):
        super().setUp()
        self.unused_port = unused_port()
        result = subprocess.run(
            [shutil.which("ctf"), "start", "--port", str(self.unused_port)]
        )
        self.assertEqual(result.returncode, 0)

    def tearDown(self):
        result = subprocess.run([shutil.which("ctf"), "stop"])
        self.assertEqual(result.returncode, 0)
        super().tearDown()