morethantext/release_tests/support/cluster.py

"""Create a cluster."""

from aiohttp import ClientSession
from release_tests.support.mttserver import MTTServer
from release_tests.support.translate import Translate


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

    def __init__(self, num=2, transurl=None, transreplies=None):
        """initialization"""
        self.translate = Translate(url=transurl, replies=transreplies)
        self.session = None
        self.servers = []
        for _ in range(num):
            self.servers.append(MTTServer())

    async def start(self):
        """Start the cluster."""
        await self.translate.start()
        for server in self.servers:
            await server.start()
        self.session = ClientSession()

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

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