Class MockJobManager

tests/autoalloc/mock/mock.py:53–92

Creates a mock job manager that intercepts job manager (PBS/Slurm) in the given `hq_env` and passes them to the given `handler`.

class MockJobManager:
    """
    Creates a mock job manager that intercepts job manager (PBS/Slurm) in the given `hq_env` and
    passes them to the given `handler`.
    """

    def __init__(
        self,
        hq_env: HqEnv,
        handler: CommandHandler,
        mocked_commands=("qsub", "qstat", "qdel", "sbatch", "scontrol", "scancel"),
    ):
        self.handler: CommandHandler = handler

        # Make sure that we're not communicating with a different test instance that could target
        # the same port.
        self.key = generate_key()

        # Prepare redirector path
        self.redirector_path = hq_env.mock.directory / "redirector.py"
        self.mocked_commands = mocked_commands
        self.hq_env = hq_env

        self.bg_server: Optional[BackgroundServer] = None

    def __enter__(self):
        assert self.bg_server is None
        self.bg_server = BackgroundServer(handler=self.handler, key=self.key)
        prepare_redirector(self.redirector_path, port=self.bg_server.port, key=self.key)
        # Link selected commands to redirector
        for cmd in self.mocked_commands:
            self.hq_env.mock.redirect_program_to_binary(cmd, self.redirector_path)

        logging.debug(f"Wrote redirector to {self.redirector_path}")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        assert self.bg_server is not None
        self.bg_server.stop()
        self.bg_server = None


def generate_key() -> str: