morethantext/release_tests/support/cluster.py

"""Create a cluster."""

from aiohttp import ClientSession
from release_tests.support import SESSION_KEY, get_port
from release_tests.support.mttserver import MTTServer
from release_tests.support.translate import Translate

ADDR = "127.56.0.1"


class Cluster:
    """cluster of MoreThanText."""

    def __init__(self, num=2, transurl=None, transreplies=None):
        """initialization"""
        self.num = num
        self.translate = Translate(url=transurl, replies=transreplies)
        self.ports = []
        self.servers = []
        self.cookies = {}

    async def start(self):
        """Start the cluster."""
        await self.translate.start()
        for _ in range(self.num):
            port = get_port()
            server = MTTServer("-a", ADDR, "-p", port)
            await server.start()
            self.servers.append(server)

    async def start_a_server(self, *args):
        """Creates a single server."""
        self.num = 1
        server = MTTServer(*args)
        await server.start()
        self.servers.append(server)

    async def start_a_cluster(self):
        """Start a service cluster."""
        for _ in range(self.num):
            port = get_port()
            while port in self.ports:
                port = get_port()
            self.ports.append(port)
            server = MTTServer("-a", ADDR, "-p", port)
            await server.start()
            self.servers.append(server)

    async def stop(self):
        """Stops cluster."""
        await self.translate.stop()
        for server in self.servers:
            await server.stop()

    async def cleanup(self):
        """Cleans up the cluster."""
        await self.stop()
        for server in self.servers:
            await server.cleanup()
            await server.stop()

    async def get(self, url):
        """get from each server in the cluster."""
        replies = []
        for server in self.servers:
            async with ClientSession() as session:
                async with session.get(
                    f"{server.baseurl}{url}", cookies=self.cookies
                ) as resp:
                    replies.append(resp)
        return replies

    async def set_session(self):
        """Sets up the session ids for the requests."""
        results = await self.get("/")
        self.cookies[SESSION_KEY] = results[0].cookies[SESSION_KEY].value