pw_rpc: Update TestMethodContext to new Method API

This change renames the TestMethodContext to NanopbTestMethod context,
and updates it to work with the new codegen API which finds a method by
ID instead of its member function pointer.

Change-Id: I0da144d5edf8990322489cc616b0f749455e4717
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/22044
Reviewed-by: Wyatt Hepler <hepler@google.com>

10 files changed, 77 insertions, 57 deletions

diff --git a/pw_rpc/BUILD b/pw_rpc/BUILD
index 3aeeb15f4..90115524d 100644
--- a/pw_rpc/BUILD
+++ b/pw_rpc/BUILD
@@ -109,12 +109,12 @@ filegroup(
         "nanopb/nanopb_method_test.cc",
         "nanopb/nanopb_method_union_test.cc",
         "nanopb/public/pw_rpc/echo_service_nanopb.h",
-        "nanopb/public/pw_rpc/nanopb_client_call.h",
         "nanopb/public/pw_rpc/internal/nanopb_common.h",
         "nanopb/public/pw_rpc/internal/nanopb_method.h",
         "nanopb/public/pw_rpc/internal/nanopb_method_union.h",
         "nanopb/public/pw_rpc/internal/service_method_traits.h",
-        "nanopb/public/pw_rpc/test_method_context.h",
+        "nanopb/public/pw_rpc/nanopb_client_call.h",
+        "nanopb/public/pw_rpc/nanopb_test_method_context.h",
         "nanopb/pw_rpc_nanopb_private/internal_test_utils.h",
         "nanopb/service_method_traits_test.cc",
         "nanopb/test.pb.c",
diff --git a/pw_rpc/nanopb/BUILD.gn b/pw_rpc/nanopb/BUILD.gn
index 9d2de1e71..452c0d542 100644
--- a/pw_rpc/nanopb/BUILD.gn
+++ b/pw_rpc/nanopb/BUILD.gn
@@ -76,7 +76,7 @@ pw_source_set("service_method_traits") {
 
 pw_source_set("test_method_context") {
   public_configs = [ ":public" ]
-  public = [ "public/pw_rpc/test_method_context.h" ]
+  public = [ "public/pw_rpc/nanopb_test_method_context.h" ]
   public_deps = [
     ":service_method_traits",
     "..:server",
@@ -106,10 +106,8 @@ pw_doc_group("docs") {
 pw_test_group("tests") {
   tests = [
     ":client_call_test",
-
-    # TODO(frolv): Re-enable these when TestMethodContext is updated.
-    # ":codegen_test",
-    # ":echo_service_test",
+    ":codegen_test",
+    ":echo_service_test",
     ":nanopb_method_test",
     ":nanopb_method_union_test",
     ":service_method_traits_test",
diff --git a/pw_rpc/nanopb/codegen_test.cc b/pw_rpc/nanopb/codegen_test.cc
index 10ab42a79..d7e96b930 100644
--- a/pw_rpc/nanopb/codegen_test.cc
+++ b/pw_rpc/nanopb/codegen_test.cc
@@ -14,7 +14,7 @@
 
 #include "gtest/gtest.h"
 #include "pw_rpc/internal/hash.h"
-#include "pw_rpc/test_method_context.h"
+#include "pw_rpc/nanopb_test_method_context.h"
 #include "pw_rpc_test_protos/test.rpc.pb.h"
 
 namespace pw::rpc {
@@ -51,7 +51,7 @@ TEST(NanopbCodegen, CompilesProperly) {
 }
 
 TEST(NanopbCodegen, InvokeUnaryRpc) {
-  TestMethodContext<&test::TestService::TestRpc> context;
+  PW_NANOPB_TEST_METHOD_CONTEXT(test::TestService, TestRpc) context;
 
   EXPECT_EQ(Status::Ok(),
             context.call({.integer = 123, .status_code = Status::Ok()}));
@@ -65,7 +65,7 @@ TEST(NanopbCodegen, InvokeUnaryRpc) {
 }
 
 TEST(NanopbCodegen, InvokeStreamingRpc) {
-  TestMethodContext<&test::TestService::TestStreamRpc> context;
+  PW_NANOPB_TEST_METHOD_CONTEXT(test::TestService, TestStreamRpc) context;
 
   context.call({.integer = 0, .status_code = Status::Aborted()});
 
@@ -87,7 +87,7 @@ TEST(NanopbCodegen, InvokeStreamingRpc) {
 }
 
 TEST(NanopbCodegen, InvokeStreamingRpc_ContextKeepsFixedNumberOfResponses) {
-  TestMethodContext<&test::TestService::TestStreamRpc, 3> context;
+  PW_NANOPB_TEST_METHOD_CONTEXT(test::TestService, TestStreamRpc, 3) context;
 
   ASSERT_EQ(3u, context.responses().max_size());
 
@@ -102,7 +102,7 @@ TEST(NanopbCodegen, InvokeStreamingRpc_ContextKeepsFixedNumberOfResponses) {
 }
 
 TEST(NanopbCodegen, InvokeStreamingRpc_ManualWriting) {
-  TestMethodContext<&test::TestService::TestStreamRpc, 3> context;
+  PW_NANOPB_TEST_METHOD_CONTEXT(test::TestService, TestStreamRpc, 3) context;
 
   ASSERT_EQ(3u, context.responses().max_size());
 
diff --git a/pw_rpc/nanopb/echo_service_test.cc b/pw_rpc/nanopb/echo_service_test.cc
index 2d4b6f231..31fbb7f37 100644
--- a/pw_rpc/nanopb/echo_service_test.cc
+++ b/pw_rpc/nanopb/echo_service_test.cc
@@ -14,19 +14,19 @@
 
 #include "gtest/gtest.h"
 #include "pw_rpc/echo_service_nanopb.h"
-#include "pw_rpc/test_method_context.h"
+#include "pw_rpc/nanopb_test_method_context.h"
 
 namespace pw::rpc {
 namespace {
 
 TEST(EchoService, Echo_EchoesRequestMessage) {
-  TestMethodContext<&EchoService::Echo> context;
+  PW_NANOPB_TEST_METHOD_CONTEXT(EchoService, Echo) context;
   ASSERT_EQ(context.call(_pw_rpc_EchoMessage{"Hello, world"}), Status::Ok());
   EXPECT_STREQ(context.response().msg, "Hello, world");
 }
 
 TEST(EchoService, Echo_EmptyRequest) {
-  TestMethodContext<&EchoService::Echo> context;
+  PW_NANOPB_TEST_METHOD_CONTEXT(EchoService, Echo) context;
   ASSERT_EQ(context.call({.msg = {}}), Status::Ok());
   EXPECT_STREQ(context.response().msg, "");
 }
diff --git a/pw_rpc/nanopb/public/pw_rpc/internal/nanopb_method_union.h b/pw_rpc/nanopb/public/pw_rpc/internal/nanopb_method_union.h
index 8997f8d92..eef6e4703 100644
--- a/pw_rpc/nanopb/public/pw_rpc/internal/nanopb_method_union.h
+++ b/pw_rpc/nanopb/public/pw_rpc/internal/nanopb_method_union.h
@@ -104,8 +104,6 @@ constexpr NanopbMethod GetNanopbMethodFor(
         id, request_fields, response_fields);
   }
 
-  // This will never be reached, as a static_assert will fail if none of the
-  // above branches are taken.
   constexpr auto fake_invoker =
       +[](ServerCall&, const int&, ServerWriter<int>&) {};
   return NanopbMethod::ServerStreaming<fake_invoker>(0, nullptr, nullptr);
diff --git a/pw_rpc/nanopb/public/pw_rpc/internal/service_method_traits.h b/pw_rpc/nanopb/public/pw_rpc/internal/service_method_traits.h
index cf5dffe3f..749789485 100644
--- a/pw_rpc/nanopb/public/pw_rpc/internal/service_method_traits.h
+++ b/pw_rpc/nanopb/public/pw_rpc/internal/service_method_traits.h
@@ -26,11 +26,14 @@ T BaseFromMember(U T::*);
 // to a member function of the service implementation to identify the service
 // class, generated service class, and Method object. This class is friended by
 // the generated service classes to give it access to the internal method list.
-template <typename Service, uint32_t method_id>
+template <auto impl_method, uint32_t method_id>
 class ServiceMethodTraits {
  public:
   ServiceMethodTraits() = delete;
 
+  // Type of the service implementation derived class.
+  using Service = typename internal::RpcTraits<decltype(impl_method)>::Service;
+
   // Type of the generic service base class.
   using BaseService =
       decltype(BaseFromMember(&Service::_PwRpcInternalGeneratedBase));
diff --git a/pw_rpc/nanopb/public/pw_rpc/test_method_context.h b/pw_rpc/nanopb/public/pw_rpc/nanopb_test_method_context.h
index 246d9665a..a0a377bff 100644
--- a/pw_rpc/nanopb/public/pw_rpc/test_method_context.h
+++ b/pw_rpc/nanopb/public/pw_rpc/nanopb_test_method_context.h
@@ -29,13 +29,13 @@
 namespace pw::rpc {
 
 // Declares a context object that may be used to invoke an RPC. The context is
-// declared with a pointer to the service member function (&Service::Method).
+// declared with the name of the implemented service and the method to invoke.
 // The RPC can then be invoked with the call method.
 //
 // For a unary RPC, context.call(request) returns the status, and the response
 // struct can be accessed via context.response().
 //
-//   pw::rpc::TestMethodContext<&my::CoolService::TheMethod> context;
+//   PW_NANOPB_TEST_METHOD_CONTEXT(my::CoolService, TheMethod) context;
 //   EXPECT_EQ(Status::Ok(), context.call({.some_arg = 123}));
 //   EXPECT_EQ(500, context.response().some_response_value);
 //
@@ -43,7 +43,7 @@ namespace pw::rpc {
 // normal RPC, the method completes when the ServerWriter's Finish method is
 // called (or it goes out of scope).
 //
-//   pw::rpc::TestMethodContext<&my::CoolService::TheStreamingMethod> context;
+//   PW_NANOPB_TEST_METHOD_CONTEXT(my::CoolService, TheStreamingMethod) context;
 //   context.call({.some_arg = 123});
 //
 //   EXPECT_TRUE(context.done());  // Check that the RPC completed
@@ -56,25 +56,33 @@ namespace pw::rpc {
 //     // iterate over the responses
 //   }
 //
-// TestMethodContext forwards its constructor arguments to the underlying
-// serivce. For example:
+// PW_NANOPB_TEST_METHOD_CONTEXT forwards its constructor arguments to the
+// underlying serivce. For example:
 //
-//   pw::rpc::TestMethodContext<&MyService::Go> context(serivce, args);
+//   PW_NANOPB_TEST_METHOD_CONTEXT(MyService, Go) context(serivce, args);
 //
-// pw::rpc::TestMethodContext takes two optional template arguments:
+// PW_NANOPB_TEST_METHOD_CONTEXT takes two optional arguments:
 //
 //   size_t max_responses: maximum responses to store; ignored unless streaming
 //   size_t output_size_bytes: buffer size; must be large enough for a packet
 //
 // Example:
 //
-//   pw::rpc::TestMethodContext<&MyService::BestMethod, 3, 256> context;
+//   PW_NANOPB_TEST_METHOD_CONTEXT(MyService, BestMethod, 3, 256) context;
 //   ASSERT_EQ(3u, context.responses().max_size());
 //
-template <auto method, size_t max_responses = 4, size_t output_size_bytes = 128>
-class TestMethodContext;
 
-// Internal classes that implement TestMethodContext.
+#define PW_NANOPB_TEST_METHOD_CONTEXT(service, method, ...)              \
+  ::pw::rpc::NanopbTestMethodContext<&service::method,                   \
+                                     ::pw::rpc::internal::Hash(#method), \
+                                     ##__VA_ARGS__>
+template <auto method,
+          uint32_t method_id,
+          size_t max_responses = 4,
+          size_t output_size_bytes = 128>
+class NanopbTestMethodContext;
+
+// Internal classes that implement NanopbTestMethodContext.
 namespace internal::test {
 
 // A ChannelOutput implementation that stores the outgoing payloads and status.
@@ -114,27 +122,32 @@ class MessageOutput final : public ChannelOutput {
 };
 
 // Collects everything needed to invoke a particular RPC.
-template <auto method, size_t max_responses, size_t output_size>
+template <auto method,
+          uint32_t method_id,
+          size_t max_responses,
+          size_t output_size>
 struct InvocationContext {
   using Request = internal::Request<method>;
   using Response = internal::Response<method>;
 
   template <typename... Args>
   InvocationContext(Args&&... args)
-      : output(ServiceMethodTraits<method>::method(), responses, buffer),
+      : output(ServiceMethodTraits<method, method_id>::method(),
+               responses,
+               buffer),
         channel(Channel::Create<123>(&output)),
         server(std::span(&channel, 1)),
         service(std::forward<Args>(args)...),
         call(static_cast<internal::Server&>(server),
              static_cast<internal::Channel&>(channel),
              service,
-             ServiceMethodTraits<method>::method()) {}
+             ServiceMethodTraits<method, method_id>::method()) {}
 
   MessageOutput<Response> output;
 
   rpc::Channel channel;
   rpc::Server server;
-  typename ServiceMethodTraits<method>::Service service;
+  typename ServiceMethodTraits<method, method_id>::Service service;
   Vector<Response, max_responses> responses;
   std::array<std::byte, output_size> buffer = {};
 
@@ -143,10 +156,10 @@ struct InvocationContext {
 
 // Method invocation context for a unary RPC. Returns the status in call() and
 // provides the response through the response() method.
-template <auto method, size_t output_size>
+template <auto method, uint32_t method_id, size_t output_size>
 class UnaryContext {
  private:
-  InvocationContext<method, 1, output_size> ctx_;
+  InvocationContext<method, method_id, 1, output_size> ctx_;
 
  public:
   using Request = typename decltype(ctx_)::Request;
@@ -172,10 +185,13 @@ class UnaryContext {
 };
 
 // Method invocation context for a server streaming RPC.
-template <auto method, size_t max_responses, size_t output_size>
+template <auto method,
+          uint32_t method_id,
+          size_t max_responses,
+          size_t output_size>
 class ServerStreamingContext {
  private:
-  InvocationContext<method, max_responses, output_size> ctx_;
+  InvocationContext<method, method_id, max_responses, output_size> ctx_;
 
  public:
   using Request = typename decltype(ctx_)::Request;
@@ -223,12 +239,13 @@ class ServerStreamingContext {
 
 // Alias to select the type of the context object to use based on which type of
 // RPC it is for.
-template <auto method, size_t responses, size_t output_size>
+template <auto method, uint32_t method_id, size_t responses, size_t output_size>
 using Context = std::tuple_element_t<
     static_cast<size_t>(internal::RpcTraits<decltype(method)>::kType),
     std::tuple<
-        internal::test::UnaryContext<method, output_size>,
-        internal::test::ServerStreamingContext<method, responses, output_size>
+        internal::test::UnaryContext<method, method_id, output_size>,
+        internal::test::
+            ServerStreamingContext<method, method_id, responses, output_size>
         // TODO(hepler): Support client and bidi streaming
         >>;
 
@@ -250,8 +267,9 @@ Status MessageOutput<Response>::SendAndReleaseBuffer(size_t size) {
 
   Result<internal::Packet> result =
       internal::Packet::FromBuffer(std::span(buffer_.data(), size));
+  PW_CHECK(result.ok());
 
-  last_status_ = result.status();
+  last_status_ = result.value().status();
 
   switch (result.value().type()) {
     case internal::PacketType::RESPONSE:
@@ -273,15 +291,20 @@ Status MessageOutput<Response>::SendAndReleaseBuffer(size_t size) {
 
 }  // namespace internal::test
 
-template <auto method, size_t max_responses, size_t output_size_bytes>
-class TestMethodContext
-    : public internal::test::Context<method, max_responses, output_size_bytes> {
+template <auto method,
+          uint32_t method_id,
+          size_t max_responses,
+          size_t output_size_bytes>
+class NanopbTestMethodContext
+    : public internal::test::
+          Context<method, method_id, max_responses, output_size_bytes> {
  public:
   // Forwards constructor arguments to the service class.
   template <typename... ServiceArgs>
-  TestMethodContext(ServiceArgs&&... service_args)
-      : internal::test::Context<method, max_responses, output_size_bytes>(
-            std::forward<ServiceArgs>(service_args)...) {}
+  NanopbTestMethodContext(ServiceArgs&&... service_args)
+      : internal::test::
+            Context<method, method_id, max_responses, output_size_bytes>(
+                std::forward<ServiceArgs>(service_args)...) {}
 };
 
 }  // namespace pw::rpc
diff --git a/pw_rpc/nanopb/service_method_traits_test.cc b/pw_rpc/nanopb/service_method_traits_test.cc
index ce5f88bce..3753bd75f 100644
--- a/pw_rpc/nanopb/service_method_traits_test.cc
+++ b/pw_rpc/nanopb/service_method_traits_test.cc
@@ -22,14 +22,14 @@
 namespace pw::rpc::internal {
 namespace {
 
-static_assert(
-    std::is_same_v<ServiceMethodTraits<EchoService, Hash("Echo")>::BaseService,
-                   generated::EchoService<EchoService>>);
+static_assert(std::is_same_v<ServiceMethodTraits<&EchoService::Echo,
+                                                 Hash("Echo")>::BaseService,
+                             generated::EchoService<EchoService>>);
 
 static_assert(
-    std::is_same_v<
-        decltype(ServiceMethodTraits<EchoService, Hash("Echo")>::method()),
-        const NanopbMethod&>);
+    std::is_same_v<decltype(ServiceMethodTraits<&EchoService::Echo,
+                                                Hash("Echo")>::method()),
+                   const NanopbMethod&>);
 
 }  // namespace
 }  // namespace pw::rpc::internal
diff --git a/pw_rpc/py/pw_rpc/codegen_nanopb.py b/pw_rpc/py/pw_rpc/codegen_nanopb.py
index 3ef422a0e..4ac3d02d4 100644
--- a/pw_rpc/py/pw_rpc/codegen_nanopb.py
+++ b/pw_rpc/py/pw_rpc/codegen_nanopb.py
@@ -156,7 +156,7 @@ def _generate_code_for_service(service: ProtoService, root: ProtoNode,
         _generate_method_lookup_function(output)
 
         output.write_line()
-        output.write_line('template <typename, uint32_t>')
+        output.write_line('template <auto, uint32_t>')
         output.write_line(
             'friend class ::pw::rpc::internal::ServiceMethodTraits;')
 
@@ -190,7 +190,7 @@ def generate_code_for_package(file_descriptor_proto, package: ProtoNode,
     output.write_line(f'#include "{nanopb_header}"\n')
 
     output.write_line('namespace pw::rpc::internal {\n')
-    output.write_line('template <typename, uint32_t>')
+    output.write_line('template <auto, uint32_t>')
     output.write_line('class ServiceMethodTraits;')
     output.write_line('\n}  // namespace pw::rpc::internal\n')
 
diff --git a/pw_rpc/raw/public/pw_rpc/internal/raw_method_union.h b/pw_rpc/raw/public/pw_rpc/internal/raw_method_union.h
index dede72026..90267c5d5 100644
--- a/pw_rpc/raw/public/pw_rpc/internal/raw_method_union.h
+++ b/pw_rpc/raw/public/pw_rpc/internal/raw_method_union.h
@@ -85,8 +85,6 @@ constexpr RawMethod GetRawMethodFor(uint32_t id) {
     return RawMethod::ServerStreaming<invoker>(id);
   }
 
-  // This will never be reached, as a static_assert will fail if none of the
-  // above branches are taken.
   constexpr auto fake_invoker =
       +[](ServerCall&, ConstByteSpan, RawServerWriter&) {};
   return RawMethod::ServerStreaming<fake_invoker>(0);