aspect/test/entry.py

#!/usr/bin/env python3

import os
import sys
import time

import flexitest

import rpc

class ServerFactory(flexitest.Factory):
    def __init__(self, port_range: list[int]):
        super().__init__(port_range)

    def create_server(self, name: str) -> flexitest.Service:
        raise NotImplementedError()

class ProcServerFactory(ServerFactory):
    def __init__(self, port_range: list[int]):
        super().__init__(port_range)

    @flexitest.with_ectx("ectx")
    def create_server(self, name: str, ectx: flexitest.EnvContext) -> flexitest.Service:
        rpc_port = self.next_port()
        rpc_url = "ws://localhost:%s" % rpc_port
        p2p_port = self.next_port()
        dd = ectx.make_service_dir(name)
        sqlite_url = "sqlite://" + os.path.realpath(os.path.join(dd, "db.sqlite")) + "?mode=rwc"
        sled_path = os.path.join(dd, "sled")
        logfile_path = os.path.join(dd, "service.log")

        cmd = ["aspect-homeserver", "-d", sqlite_url, "-s", sled_path, "-r", str(rpc_port), "-p", str(p2p_port)]
        props = {
            "rpc_port": rpc_port,
            "rpc_url": rpc_url,
            "p2p_port": p2p_port,
            "p2p_addr": "127.0.0.1:%s" % p2p_port,
            "logfile": logfile_path
        }

        svc = flexitest.service.ProcService(props, cmd, stdout=logfile_path)
        _inject_rpc(svc)
        return svc

class ClientFactory(flexitest.Factory):
    def __init__(self, port_range: list[int]):
        super().__init__(port_range)

    def create_client(self, name: str) -> flexitest.Service:
        raise NotImplementedError()

class ProcClientFactory(ClientFactory):
    def __init__(self, port_range: list[int]):
        super().__init__(port_range)

    @flexitest.with_ectx("ectx")
    def create_client(self, name: str, ectx: flexitest.EnvContext) -> flexitest.Service:
        rpc_port = self.next_port()
        rpc_url = "ws://localhost:%s" % rpc_port
        dd = ectx.make_service_dir(name)
        sqlite_url = "sqlite://" + os.path.realpath(os.path.join(dd, "db.sqlite")) + "?mode=rwc"
        logfile_path = os.path.join(dd, "service.log")

        cmd = ["aspect-user", "-d", dd, "-u", sqlite_url, "-r", str(rpc_port)]
        props = {
            "rpc_port": rpc_port,
            "rpc_url": rpc_url,
            "logfile": logfile_path
        }

        svc = flexitest.service.ProcService(props, cmd, stdout=logfile_path)
        _inject_rpc(svc)
        return svc

class BareEnvConfig(flexitest.EnvConfig):
    """Environment that just inits some number of clients and servers."""
    def __init__(self, nserv, nclients):
        self.num_servers = nserv
        self.num_clients = nclients

    def init(self, ectx: flexitest.EnvContext) -> flexitest.LiveEnv:
        sfac = ectx.get_factory("server")
        cfac = ectx.get_factory("client")
        svcs = {}

        for i in range(self.num_servers):
            name = "serv%s" % i
            svc = sfac.create_server(name)
            svcs[name] = svc
            svc.start()

        for i in range(self.num_clients):
            name = "client%s" % i
            svc = cfac.create_client(name)
            svcs[name] = svc
            svc.start()

        time.sleep(1)

        return flexitest.LiveEnv(svcs)

class SimpleEnvConfig(flexitest.EnvConfig):
    """
    Environment that creates a single server and some number of clients that are
    all connected and registered to it.
    """
    def __init__(self, nclients):
        self.num_clients = nclients

    def init(self, ectx: flexitest.EnvContext) -> flexitest.LiveEnv:
        sfac = ectx.get_factory("server")
        cfac = ectx.get_factory("client")
        svcs = {}

        s = sfac.create_server("serv")
        svcs["serv"] = s
        serv_addr = s.get_prop("p2p_addr")
        time.sleep(1)

        for i in range(self.num_clients):
            name = "cli%s" % i
            print("creating client", name)
            c = cfac.create_client(name)
            c.start()
            time.sleep(1)
            svcs[name] = c
            crpc = c.create_rpc()
            crpc.aspc_register(serv_addr, name, name)

        return flexitest.LiveEnv(svcs)

def _inject_rpc(svc: flexitest.Service):
    def _create_rpc():
        return rpc.RpcClient(svc.get_prop("rpc_url"))
    setattr(svc, "create_rpc", _create_rpc)

def main(argv):
    test_dir = os.path.dirname(os.path.abspath(__file__))

    # Create datadir.
    datadir_root = flexitest.create_datadir_in_workspace(os.path.join(test_dir, "_dd"))

    # Probe tests.
    modules = flexitest.runtime.scan_dir_for_modules(test_dir)
    tests = flexitest.runtime.load_candidate_modules(modules)
    print(tests)

    # Init factors to be able to create the test envs with.
    sfac = ProcServerFactory([12300 + i for i in range(20)])
    cfac = ProcClientFactory([12400 + i for i in range(20)])
    factories = {"server": sfac, "client": cfac}

    # Prepare env configs.
    bare_env = BareEnvConfig(1, 2)
    basic_env = SimpleEnvConfig(1)
    basic2_env = SimpleEnvConfig(2)
    env_configs = {"bare": bare_env, "basic": basic_env, "basic2": basic2_env}

    # Set up the runtime and prepare tests.
    rt = flexitest.TestRuntime(env_configs, datadir_root, factories)
    rt.prepare_registered_tests()

    # Run the tests and then dump the results.
    arg_test_names = argv[1:]
    if len(arg_test_names) > 0:
        tests = arg_test_names

    results = rt.run_tests(tests)
    rt.save_json_file("results.json", results)
    flexitest.dump_results(results)
    flexitest.fail_on_error(results)

    return 0

if __name__ == "__main__":
    sys.exit(main(sys.argv))