atlas-iac/testing/tests/test_hermes_scm_broker_support.py

"""Shared fixtures for credential-isolated Hermes SCM broker tests."""

from __future__ import annotations

import hashlib
import importlib.util
import io
import sys
import zlib
from email.message import Message
from pathlib import Path

ROOT = Path(__file__).parents[2]
SCRIPTS = ROOT / "services/hermes/scm-common/scripts"
sys.path.insert(0, str(SCRIPTS))


def _load(name: str):
    path = SCRIPTS / f"{name}.py"
    spec = importlib.util.spec_from_file_location(f"test_{name}", path)
    assert spec and spec.loader
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def _load_path(name: str, path: Path):
    spec = importlib.util.spec_from_file_location(name, path)
    assert spec and spec.loader
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


class Response:
    def __init__(self, body: bytes, *, status: int = 200, content_type: str):
        self.body = body
        self.stream = io.BytesIO(body)
        self.status = status
        self.headers = Message()
        self.headers["Content-Type"] = content_type

    def __enter__(self):
        return self

    def __exit__(self, *_args):
        return False

    def read(self, limit=-1):
        return self.stream.read(limit)


def _object_entry(type_code: int, payload: bytes) -> bytes:
    """Encode one non-delta pack entry with a real header and zlib body."""
    size = len(payload)
    byte = (type_code << 4) | (size & 0x0F)
    size >>= 4
    header = bytearray()
    while size:
        header.append(byte | 0x80)
        byte = size & 0x7F
        size >>= 7
    header.append(byte)
    return bytes(header) + zlib.compress(payload)


def _pack_of(entries: list[bytes]) -> bytes:
    """Assemble encoded entries into a checksummed version-2 pack."""
    body = b"PACK" + (2).to_bytes(4, "big") + len(entries).to_bytes(4, "big")
    for entry in entries:
        body += entry
    return body + hashlib.sha1(body).digest()


def _blob_pack(*payloads: bytes) -> bytes:
    return _pack_of([_object_entry(3, payload) for payload in payloads])


def _receive_command(old: bytes, new: bytes, ref: bytes, pack: bytes = b"") -> bytes:
    command = old + b" " + new + b" " + ref + b"\x00report-status\n"
    framed = f"{len(command) + 4:04x}".encode() + command + b"0000"
    return framed + (pack or _pack_of([]))