chattolib/tests/test_protoc_plugin.py

"""The plugin turns a CodeGeneratorRequest into async+sync client stubs."""

import importlib.util
import os

from chattolib._protoc_plugin import generate

_CODEGEN_DIR = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
    "src",
    "chattolib",
    "_codegen",
)


def _load(name):
    path = os.path.join(_CODEGEN_DIR, f"{name}.py")
    spec = importlib.util.spec_from_file_location(f"_test_codegen_{name}", path)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)  # type: ignore[union-attr]
    return mod


g = _load("protoc_gen_request_pb2")


def _request_with_room_service() -> g.CodeGeneratorRequest:
    req = g.CodeGeneratorRequest()
    req.file_to_generate.append("chatto/api/v1/rooms.proto")
    fd = req.proto_file.add()
    fd.name = "chatto/api/v1/rooms.proto"
    fd.package = "chatto.api.v1"
    svc = fd.service.add()
    svc.name = "RoomService"
    for name in ("JoinRoom", "CreateRoom"):
        m = svc.method.add()
        m.name = name
        m.input_type = f".chatto.api.v1.{name}Request"
        m.output_type = f".chatto.api.v1.{name}Response"
    return req


def test_emits_one_file_per_service_with_both_clients():
    resp = generate(_request_with_room_service())
    assert len(resp.file) == 1
    out = resp.file[0]
    assert out.name.endswith("rooms_connect.py")
    content = out.content
    # Both client classes, correct base classes, one method per RPC.
    assert "class RoomServiceClient(ConnectClient):" in content
    assert "class RoomServiceClientSync(ConnectClientSync):" in content
    assert "async def join_room(" in content
    assert "async def create_room(" in content
    assert "def join_room(" in content  # sync variant
    # Imports only from chattolib._connect and the sibling pb2.
    assert "from chattolib._connect import ConnectClient, ConnectClientSync, MethodInfo" in content
    assert "import chatto.api.v1.rooms_pb2 as _pb2" in content
    # No connectrpc anywhere.
    assert "connectrpc" not in content
    # use_get is dropped.
    assert "use_get" not in content


def test_skips_files_without_services():
    req = g.CodeGeneratorRequest()
    req.file_to_generate.append("chatto/realtime/v1/realtime.proto")
    fd = req.proto_file.add()
    fd.name = "chatto/realtime/v1/realtime.proto"
    fd.package = "chatto.realtime.v1"
    fd.message_type.add().name = "RealtimeClientFrame"  # messages, no service
    resp = generate(req)
    assert len(resp.file) == 0


def _request_with_cross_file_rpc() -> g.CodeGeneratorRequest:
    """A service in account.proto whose UpdatePresence RPC uses a message that
    is defined in presence.proto — the cross-file case the plugin must handle."""
    req = g.CodeGeneratorRequest()
    req.file_to_generate.append("chatto/api/v1/account.proto")

    # account.proto: declares the service + its own UpdateProfile messages.
    acc = req.proto_file.add()
    acc.name = "chatto/api/v1/account.proto"
    acc.package = "chatto.api.v1"
    acc.message_type.add().name = "UpdateProfileRequest"
    acc.message_type.add().name = "UpdateProfileResponse"
    svc = acc.service.add()
    svc.name = "MyAccountService"
    m = svc.method.add()
    m.name = "UpdatePresence"
    m.input_type = ".chatto.api.v1.UpdatePresenceRequest"
    m.output_type = ".chatto.api.v1.UpdatePresenceResponse"
    m2 = svc.method.add()
    m2.name = "UpdateProfile"
    m2.input_type = ".chatto.api.v1.UpdateProfileRequest"
    m2.output_type = ".chatto.api.v1.UpdateProfileResponse"

    # presence.proto: defines the UpdatePresence request/response.
    pres = req.proto_file.add()
    pres.name = "chatto/api/v1/presence.proto"
    pres.package = "chatto.api.v1"
    pres.message_type.add().name = "UpdatePresenceRequest"
    pres.message_type.add().name = "UpdatePresenceResponse"
    return req


def test_cross_file_rpc_uses_the_defining_module():
    resp = generate(_request_with_cross_file_rpc())
    assert len(resp.file) == 1
    content = resp.file[0].content
    # The account_pb2 import is still there for the same-file messages.
    assert "import chatto.api.v1.account_pb2 as _pb2" in content
    # The presence messages must be imported from their own module and the
    # UpdatePresence method must reference *that* module, not _pb2.
    assert "import chatto.api.v1.presence_pb2 as _pb2_presence" in content
    assert "_pb2_presence.UpdatePresenceRequest" in content
    assert "_pb2_presence.UpdatePresenceResponse" in content
    # The same-file method keeps using the plain _pb2 alias.
    assert "_pb2.UpdateProfileRequest" in content
    # Crucially: UpdatePresenceRequest must NOT be referenced via _pb2 (the
    # original bug was `_pb2.UpdatePresenceRequest` -> AttributeError).
    assert "_pb2.UpdatePresenceRequest" not in content
    assert "_pb2.UpdatePresenceResponse" not in content