test_protoc_plugin.py 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. """The plugin turns a CodeGeneratorRequest into async+sync client stubs."""
  2. import importlib.util
  3. import os
  4. from chattolib._protoc_plugin import generate
  5. _CODEGEN_DIR = os.path.join(
  6. os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
  7. "src",
  8. "chattolib",
  9. "_codegen",
  10. )
  11. def _load(name):
  12. path = os.path.join(_CODEGEN_DIR, f"{name}.py")
  13. spec = importlib.util.spec_from_file_location(f"_test_codegen_{name}", path)
  14. mod = importlib.util.module_from_spec(spec)
  15. spec.loader.exec_module(mod) # type: ignore[union-attr]
  16. return mod
  17. g = _load("protoc_gen_request_pb2")
  18. def _request_with_room_service() -> g.CodeGeneratorRequest:
  19. req = g.CodeGeneratorRequest()
  20. req.file_to_generate.append("chatto/api/v1/rooms.proto")
  21. fd = req.proto_file.add()
  22. fd.name = "chatto/api/v1/rooms.proto"
  23. fd.package = "chatto.api.v1"
  24. svc = fd.service.add()
  25. svc.name = "RoomService"
  26. for name in ("JoinRoom", "CreateRoom"):
  27. m = svc.method.add()
  28. m.name = name
  29. m.input_type = f".chatto.api.v1.{name}Request"
  30. m.output_type = f".chatto.api.v1.{name}Response"
  31. return req
  32. def test_emits_one_file_per_service_with_both_clients():
  33. resp = generate(_request_with_room_service())
  34. assert len(resp.file) == 1
  35. out = resp.file[0]
  36. assert out.name.endswith("rooms_connect.py")
  37. content = out.content
  38. # Both client classes, correct base classes, one method per RPC.
  39. assert "class RoomServiceClient(ConnectClient):" in content
  40. assert "class RoomServiceClientSync(ConnectClientSync):" in content
  41. assert "async def join_room(" in content
  42. assert "async def create_room(" in content
  43. assert "def join_room(" in content # sync variant
  44. # Imports only from chattolib._connect and the sibling pb2.
  45. assert "from chattolib._connect import ConnectClient, ConnectClientSync, MethodInfo" in content
  46. assert "import chatto.api.v1.rooms_pb2 as _pb2" in content
  47. # No connectrpc anywhere.
  48. assert "connectrpc" not in content
  49. # use_get is dropped.
  50. assert "use_get" not in content
  51. def test_skips_files_without_services():
  52. req = g.CodeGeneratorRequest()
  53. req.file_to_generate.append("chatto/realtime/v1/realtime.proto")
  54. fd = req.proto_file.add()
  55. fd.name = "chatto/realtime/v1/realtime.proto"
  56. fd.package = "chatto.realtime.v1"
  57. fd.message_type.add().name = "RealtimeClientFrame" # messages, no service
  58. resp = generate(req)
  59. assert len(resp.file) == 0