pw_rpc: Add RawTestMethodContext

This change adds a TestMethodContext class for raw RPC methods. Its API
is identical to the nanopb one, but using byte spans for requests and
responses instead of structs.

Change-Id: Icf92f026882bcb6623478762d01bfde4e2deadc2
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/23170
Commit-Queue: Alexei Frolov <frolv@google.com>
Reviewed-by: Wyatt Hepler <hepler@google.com>
diff --git a/pw_rpc/BUILD b/pw_rpc/BUILD
index 9011552..a7c4a83 100644
--- a/pw_rpc/BUILD
+++ b/pw_rpc/BUILD
@@ -83,6 +83,15 @@
 )
 
 pw_cc_library(
+    name = "service_method_traits",
+    hdrs = [
+        "public/pw_rpc/internal/service_method_traits.h",
+    ],
+    includes = ["public"],
+    deps = [ ":server"  ],
+)
+
+pw_cc_library(
     name = "internal_test_utils",
     hdrs = [
         "public/pw_rpc/internal/test_method.h",
@@ -112,11 +121,11 @@
         "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/internal/nanopb_service_method_traits.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/nanopb_service_method_traits_test.cc",
         "nanopb/test.pb.c",
         "nanopb/test.pb.h",
         "nanopb/test_rpc.pb.h",
diff --git a/pw_rpc/BUILD.gn b/pw_rpc/BUILD.gn
index a04c2e6..01228ad 100644
--- a/pw_rpc/BUILD.gn
+++ b/pw_rpc/BUILD.gn
@@ -86,6 +86,12 @@
   friend = [ "./*" ]
 }
 
+pw_source_set("service_method_traits") {
+  public = [ "public/pw_rpc/internal/service_method_traits.h" ]
+  public_deps = [ ":server" ]
+  visibility = [ "./*" ]
+}
+
 pw_source_set("test_utils") {
   public = [
     "public/pw_rpc/internal/test_method.h",
diff --git a/pw_rpc/nanopb/BUILD.gn b/pw_rpc/nanopb/BUILD.gn
index 637a078..39edd13 100644
--- a/pw_rpc/nanopb/BUILD.gn
+++ b/pw_rpc/nanopb/BUILD.gn
@@ -67,10 +67,10 @@
 
 pw_source_set("service_method_traits") {
   public_configs = [ ":public" ]
-  public = [ "public/pw_rpc/internal/service_method_traits.h" ]
+  public = [ "public/pw_rpc/internal/nanopb_service_method_traits.h" ]
   public_deps = [
-    ":method",
-    "..:server",
+    ":method_union",
+    "..:service_method_traits",
   ]
 }
 
@@ -110,7 +110,7 @@
     ":echo_service_test",
     ":nanopb_method_test",
     ":nanopb_method_union_test",
-    ":service_method_traits_test",
+    ":nanopb_service_method_traits_test",
   ]
 }
 
@@ -171,7 +171,7 @@
   enable_if = dir_pw_third_party_nanopb != ""
 }
 
-pw_test("service_method_traits_test") {
+pw_test("nanopb_service_method_traits_test") {
   deps = [
     ":echo_service",
     ":method",
@@ -179,6 +179,6 @@
     ":test_method_context",
     "..:test_protos.nanopb_rpc",
   ]
-  sources = [ "service_method_traits_test.cc" ]
+  sources = [ "nanopb_service_method_traits_test.cc" ]
   enable_if = dir_pw_third_party_nanopb != ""
 }
diff --git a/pw_rpc/nanopb/service_method_traits_test.cc b/pw_rpc/nanopb/nanopb_service_method_traits_test.cc
similarity index 67%
rename from pw_rpc/nanopb/service_method_traits_test.cc
rename to pw_rpc/nanopb/nanopb_service_method_traits_test.cc
index 3753bd7..8043e12 100644
--- a/pw_rpc/nanopb/service_method_traits_test.cc
+++ b/pw_rpc/nanopb/nanopb_service_method_traits_test.cc
@@ -12,7 +12,7 @@
 // License for the specific language governing permissions and limitations under
 // the License.
 
-#include "pw_rpc/internal/service_method_traits.h"
+#include "pw_rpc/internal/nanopb_service_method_traits.h"
 
 #include <type_traits>
 
@@ -22,13 +22,9 @@
 namespace pw::rpc::internal {
 namespace {
 
-static_assert(std::is_same_v<ServiceMethodTraits<&EchoService::Echo,
-                                                 Hash("Echo")>::BaseService,
-                             generated::EchoService<EchoService>>);
-
 static_assert(
-    std::is_same_v<decltype(ServiceMethodTraits<&EchoService::Echo,
-                                                Hash("Echo")>::method()),
+    std::is_same_v<decltype(NanopbServiceMethodTraits<&EchoService::Echo,
+                                                      Hash("Echo")>::method()),
                    const NanopbMethod&>);
 
 }  // namespace
diff --git a/pw_rpc/nanopb/public/pw_rpc/internal/nanopb_service_method_traits.h b/pw_rpc/nanopb/public/pw_rpc/internal/nanopb_service_method_traits.h
new file mode 100644
index 0000000..866ceef
--- /dev/null
+++ b/pw_rpc/nanopb/public/pw_rpc/internal/nanopb_service_method_traits.h
@@ -0,0 +1,27 @@
+// Copyright 2020 The Pigweed Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not
+// use this file except in compliance with the License. You may obtain a copy of
+// the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations under
+// the License.
+#pragma once
+
+#include "pw_rpc/internal/nanopb_method_union.h"
+#include "pw_rpc/internal/service_method_traits.h"
+
+namespace pw::rpc::internal {
+
+template <auto impl_method, uint32_t method_id>
+using NanopbServiceMethodTraits =
+    ServiceMethodTraits<&MethodBaseService<impl_method>::NanopbMethodFor,
+                        impl_method,
+                        method_id>;
+
+}  // namespace pw::rpc::internal
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
deleted file mode 100644
index 7497894..0000000
--- a/pw_rpc/nanopb/public/pw_rpc/internal/service_method_traits.h
+++ /dev/null
@@ -1,50 +0,0 @@
-// Copyright 2020 The Pigweed Authors
-//
-// Licensed under the Apache License, Version 2.0 (the "License"); you may not
-// use this file except in compliance with the License. You may obtain a copy of
-// the License at
-//
-//     https://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
-// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
-// License for the specific language governing permissions and limitations under
-// the License.
-#pragma once
-
-#include "pw_rpc/internal/nanopb_method.h"
-
-namespace pw::rpc::internal {
-
-// Identifies a base class from a member function it defines. This should be
-// used with decltype to retrieve the base class.
-template <typename T, typename U>
-T BaseFromMember(U T::*);
-
-// Gets information about a service and method at compile-time. Uses a pointer
-// 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 <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));
-
-  // Reference to the Method object corresponding to this method.
-  static constexpr const NanopbMethod& method() {
-    return *BaseService::NanopbMethodFor(method_id);
-  }
-
-  static_assert(BaseService::NanopbMethodFor(method_id) != nullptr,
-                "The selected function is not an RPC service method");
-};
-
-}  // namespace pw::rpc::internal
diff --git a/pw_rpc/nanopb/public/pw_rpc/nanopb_test_method_context.h b/pw_rpc/nanopb/public/pw_rpc/nanopb_test_method_context.h
index a0a377b..6ff26d9 100644
--- a/pw_rpc/nanopb/public/pw_rpc/nanopb_test_method_context.h
+++ b/pw_rpc/nanopb/public/pw_rpc/nanopb_test_method_context.h
@@ -22,9 +22,9 @@
 #include "pw_rpc/channel.h"
 #include "pw_rpc/internal/hash.h"
 #include "pw_rpc/internal/nanopb_method.h"
+#include "pw_rpc/internal/nanopb_service_method_traits.h"
 #include "pw_rpc/internal/packet.h"
 #include "pw_rpc/internal/server.h"
-#include "pw_rpc/internal/service_method_traits.h"
 
 namespace pw::rpc {
 
@@ -59,7 +59,7 @@
 // PW_NANOPB_TEST_METHOD_CONTEXT forwards its constructor arguments to the
 // underlying serivce. For example:
 //
-//   PW_NANOPB_TEST_METHOD_CONTEXT(MyService, Go) context(serivce, args);
+//   PW_NANOPB_TEST_METHOD_CONTEXT(MyService, Go) context(service, args);
 //
 // PW_NANOPB_TEST_METHOD_CONTEXT takes two optional arguments:
 //
@@ -83,7 +83,7 @@
 class NanopbTestMethodContext;
 
 // Internal classes that implement NanopbTestMethodContext.
-namespace internal::test {
+namespace internal::test::nanopb {
 
 // A ChannelOutput implementation that stores the outgoing payloads and status.
 template <typename Response>
@@ -92,7 +92,7 @@
   MessageOutput(const internal::NanopbMethod& method,
                 Vector<Response>& responses,
                 std::span<std::byte> buffer)
-      : ChannelOutput("internal::test::MessageOutput"),
+      : ChannelOutput("internal::test::nanopb::MessageOutput"),
         method_(method),
         responses_(responses),
         buffer_(buffer) {
@@ -132,7 +132,7 @@
 
   template <typename... Args>
   InvocationContext(Args&&... args)
-      : output(ServiceMethodTraits<method, method_id>::method(),
+      : output(NanopbServiceMethodTraits<method, method_id>::method(),
                responses,
                buffer),
         channel(Channel::Create<123>(&output)),
@@ -141,13 +141,13 @@
         call(static_cast<internal::Server&>(server),
              static_cast<internal::Channel&>(channel),
              service,
-             ServiceMethodTraits<method, method_id>::method()) {}
+             NanopbServiceMethodTraits<method, method_id>::method()) {}
 
   MessageOutput<Response> output;
 
   rpc::Channel channel;
   rpc::Server server;
-  typename ServiceMethodTraits<method, method_id>::Service service;
+  typename NanopbServiceMethodTraits<method, method_id>::Service service;
   Vector<Response, max_responses> responses;
   std::array<std::byte, output_size> buffer = {};
 
@@ -242,12 +242,10 @@
 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, method_id, output_size>,
-        internal::test::
-            ServerStreamingContext<method, method_id, responses, output_size>
-        // TODO(hepler): Support client and bidi streaming
-        >>;
+    std::tuple<UnaryContext<method, method_id, output_size>,
+               ServerStreamingContext<method, method_id, responses, output_size>
+               // TODO(hepler): Support client and bidi streaming
+               >>;
 
 template <typename Response>
 void MessageOutput<Response>::clear() {
@@ -289,20 +287,20 @@
   return Status::Ok();
 }
 
-}  // namespace internal::test
+}  // namespace internal::test::nanopb
 
 template <auto method,
           uint32_t method_id,
           size_t max_responses,
           size_t output_size_bytes>
 class NanopbTestMethodContext
-    : public internal::test::
+    : public internal::test::nanopb::
           Context<method, method_id, max_responses, output_size_bytes> {
  public:
   // Forwards constructor arguments to the service class.
   template <typename... ServiceArgs>
   NanopbTestMethodContext(ServiceArgs&&... service_args)
-      : internal::test::
+      : internal::test::nanopb::
             Context<method, method_id, max_responses, output_size_bytes>(
                 std::forward<ServiceArgs>(service_args)...) {}
 };
diff --git a/pw_rpc/public/pw_rpc/internal/method_union.h b/pw_rpc/public/pw_rpc/internal/method_union.h
index 5eeedb4..9d491af 100644
--- a/pw_rpc/public/pw_rpc/internal/method_union.h
+++ b/pw_rpc/public/pw_rpc/internal/method_union.h
@@ -42,12 +42,29 @@
   // Specializations must set Implementation as an alias for their method
   // implementation class.
   using Implementation = Method;
+
+  // Specializations must set Service as an alias to the implemented service
+  // class.
+  using Service = rpc::Service;
 };
 
 template <auto method>
 using MethodImplementation =
     typename MethodTraits<decltype(method)>::Implementation;
 
+template <auto method>
+using MethodService = typename MethodTraits<decltype(method)>::Service;
+
+// Identifies a base class from a member function it defines. This should be
+// used with decltype to retrieve the base class.
+template <typename T, typename U>
+T BaseFromMember(U T::*);
+
+// The base generated service of an implemented RPC method.
+template <auto method>
+using MethodBaseService = decltype(
+    BaseFromMember(&MethodService<method>::_PwRpcInternalGeneratedBase));
+
 class CoreMethodUnion : public MethodUnion {
  public:
   constexpr const Method& method() const { return impl_.method; }
diff --git a/pw_rpc/public/pw_rpc/internal/service_method_traits.h b/pw_rpc/public/pw_rpc/internal/service_method_traits.h
new file mode 100644
index 0000000..0e26b8d
--- /dev/null
+++ b/pw_rpc/public/pw_rpc/internal/service_method_traits.h
@@ -0,0 +1,42 @@
+// Copyright 2020 The Pigweed Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not
+// use this file except in compliance with the License. You may obtain a copy of
+// the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations under
+// the License.
+#pragma once
+
+#include "pw_rpc/internal/method_union.h"
+
+namespace pw::rpc::internal {
+
+// Gets information about a service and method at compile-time. Uses a pointer
+// to a member function of the service implementation to identify the service
+// class, generated service class, and Method object.
+template <auto lookup_function, auto impl_method, uint32_t method_id>
+class ServiceMethodTraits {
+ public:
+  ServiceMethodTraits() = delete;
+
+  // Type of the service implementation derived class.
+  using Service = MethodService<impl_method>;
+
+  using MethodImpl = MethodImplementation<impl_method>;
+
+  // Reference to the Method object corresponding to this method.
+  static constexpr const MethodImpl& method() {
+    return *lookup_function(method_id);
+  }
+
+  static_assert(lookup_function(method_id) != nullptr,
+                "The selected function is not an RPC service method");
+};
+
+}  // 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 433829a..9e5ca78 100644
--- a/pw_rpc/py/pw_rpc/codegen_nanopb.py
+++ b/pw_rpc/py/pw_rpc/codegen_nanopb.py
@@ -122,6 +122,9 @@
         output.write_line(
             'constexpr void _PwRpcInternalGeneratedBase() const {}')
 
+        output.write_line()
+        _generate_method_lookup_function(output)
+
     service_name_hash = pw_rpc.ids.calculate(service.proto_path())
     output.write_line('\n private:')
 
@@ -142,14 +145,7 @@
             for method in service.methods():
                 _generate_method_descriptor(method, output)
 
-        output.write_line('};\n')
-
-        _generate_method_lookup_function(output)
-
-        output.write_line()
-        output.write_line('template <auto, uint32_t>')
-        output.write_line(
-            'friend class ::pw::rpc::internal::ServiceMethodTraits;')
+        output.write_line('};')
 
     output.write_line('};')
 
@@ -257,11 +253,6 @@
         file_descriptor_proto.name)
     output.write_line(f'#include "{nanopb_header}"\n')
 
-    output.write_line('namespace pw::rpc::internal {\n')
-    output.write_line('template <auto, uint32_t>')
-    output.write_line('class ServiceMethodTraits;')
-    output.write_line('\n}  // namespace pw::rpc::internal\n')
-
     if package.cpp_namespace():
         file_namespace = package.cpp_namespace()
         if file_namespace.startswith('::'):
diff --git a/pw_rpc/py/pw_rpc/codegen_raw.py b/pw_rpc/py/pw_rpc/codegen_raw.py
index e8661b4..6babe1c 100644
--- a/pw_rpc/py/pw_rpc/codegen_raw.py
+++ b/pw_rpc/py/pw_rpc/codegen_raw.py
@@ -48,6 +48,27 @@
     output.write_line(f'    0x{method_id:08x}),  // Hash of "{method.name()}"')
 
 
+def _generate_method_lookup_function(output: OutputFile):
+    """Generates a function that gets a Method object from its ID."""
+    raw_method = f'{RPC_NAMESPACE}::internal::RawMethod'
+
+    output.write_line(f'static constexpr const {raw_method}* RawMethodFor(')
+    output.write_line('    uint32_t id) {')
+
+    with output.indent():
+        output.write_line('for (auto& method : kMethods) {')
+        with output.indent():
+            output.write_line('if (method.raw_method().id() == id) {')
+            output.write_line(f'  return &static_cast<const {raw_method}&>(')
+            output.write_line('    method.raw_method());')
+            output.write_line('}')
+        output.write_line('}')
+
+        output.write_line('return nullptr;')
+
+    output.write_line('}')
+
+
 def _generate_code_for_service(service: ProtoService, root: ProtoNode,
                                output: OutputFile) -> None:
     """Generates a C++ base class for a raw RPC service."""
@@ -84,6 +105,9 @@
         output.write_line(
             'constexpr void _PwRpcInternalGeneratedBase() const {}')
 
+        output.write_line()
+        _generate_method_lookup_function(output)
+
     service_name_hash = pw_rpc.ids.calculate(service.proto_path())
     output.write_line('\n private:')
 
@@ -104,7 +128,7 @@
             for method in service.methods():
                 _generate_method_descriptor(method, output)
 
-        output.write_line('};\n')
+        output.write_line('};')
 
     output.write_line('};')
 
diff --git a/pw_rpc/raw/BUILD b/pw_rpc/raw/BUILD
index df95275..229718e 100644
--- a/pw_rpc/raw/BUILD
+++ b/pw_rpc/raw/BUILD
@@ -46,6 +46,29 @@
     ]
 )
 
+pw_cc_library(
+    name = "service_method_traits",
+    hdrs = [
+        "public/pw_rpc/internal/raw_service_method_traits.h",
+    ],
+    deps = [
+        ":method_union",
+        "//pw_rpc:service_method_traits",
+    ]
+)
+
+pw_cc_library(
+    name = "test_method_context",
+    hdrs = [
+        "public/pw_rpc/raw_test_method_context.h",
+    ],
+    deps = [
+        ":method_union",
+        "//pw_assert",
+        "//pw_containers",
+    ]
+)
+
 pw_cc_test(
     name = "codegen_test",
     srcs = [
diff --git a/pw_rpc/raw/BUILD.gn b/pw_rpc/raw/BUILD.gn
index 21c88c5..89d0492 100644
--- a/pw_rpc/raw/BUILD.gn
+++ b/pw_rpc/raw/BUILD.gn
@@ -40,6 +40,25 @@
   public_deps = [ ":method" ]
 }
 
+pw_source_set("service_method_traits") {
+  public_configs = [ ":public" ]
+  public = [ "public/pw_rpc/internal/raw_service_method_traits.h" ]
+  public_deps = [
+    ":method_union",
+    "..:service_method_traits",
+  ]
+}
+
+pw_source_set("test_method_context") {
+  public_configs = [ ":public" ]
+  public = [ "public/pw_rpc/raw_test_method_context.h" ]
+  public_deps = [
+    ":service_method_traits",
+    dir_pw_assert,
+    dir_pw_containers,
+  ]
+}
+
 pw_test_group("tests") {
   tests = [
     ":codegen_test",
@@ -50,6 +69,7 @@
 
 pw_test("codegen_test") {
   deps = [
+    ":test_method_context",
     "..:test_protos.pwpb",
     "..:test_protos.raw_rpc",
     dir_pw_protobuf,
diff --git a/pw_rpc/raw/codegen_test.cc b/pw_rpc/raw/codegen_test.cc
index a9fcd59..7041cc1 100644
--- a/pw_rpc/raw/codegen_test.cc
+++ b/pw_rpc/raw/codegen_test.cc
@@ -15,6 +15,7 @@
 #include "gtest/gtest.h"
 #include "pw_protobuf/decoder.h"
 #include "pw_rpc/internal/hash.h"
+#include "pw_rpc/raw_test_method_context.h"
 #include "pw_rpc_test_protos/test.pwpb.h"
 #include "pw_rpc_test_protos/test.raw_rpc.pb.h"
 
@@ -27,7 +28,36 @@
                          ConstByteSpan request,
                          ByteSpan response) {
     int64_t integer;
-    uint32_t status_code;
+    Status status;
+    DecodeRequest(request, integer, status);
+
+    protobuf::NestedEncoder encoder(response);
+    TestResponse::Encoder test_response(&encoder);
+    test_response.WriteValue(integer + 1);
+
+    return StatusWithSize(status, encoder.Encode().value().size());
+  }
+
+  void TestStreamRpc(ServerContext&,
+                     ConstByteSpan request,
+                     RawServerWriter& writer) {
+    int64_t integer;
+    Status status;
+    DecodeRequest(request, integer, status);
+
+    for (int i = 0; i < integer; ++i) {
+      ByteSpan buffer = writer.PayloadBuffer();
+      protobuf::NestedEncoder encoder(buffer);
+      TestStreamResponse::Encoder test_stream_response(&encoder);
+      test_stream_response.WriteNumber(i);
+      writer.Write(encoder.Encode().value());
+    }
+
+    writer.Finish(status);
+  }
+
+ private:
+  void DecodeRequest(ConstByteSpan request, int64_t& integer, Status& status) {
     protobuf::Decoder decoder(request);
 
     while (decoder.Next().ok()) {
@@ -35,21 +65,15 @@
         case TestRequest::Fields::INTEGER:
           decoder.ReadInt64(&integer);
           break;
-        case TestRequest::Fields::STATUS_CODE:
+        case TestRequest::Fields::STATUS_CODE: {
+          uint32_t status_code;
           decoder.ReadUint32(&status_code);
+          status = static_cast<Status::Code>(status_code);
           break;
+        }
       }
     }
-
-    protobuf::NestedEncoder encoder(response);
-    TestResponse::Encoder test_response(&encoder);
-    test_response.WriteValue(integer + 1);
-
-    return StatusWithSize(static_cast<Status::Code>(status_code),
-                          encoder.Encode().value().size());
   }
-
-  void TestStreamRpc(ServerContext&, ConstByteSpan, RawServerWriter&) {}
 };
 
 }  // namespace test
@@ -62,5 +86,62 @@
   EXPECT_STREQ(service.name(), "TestService");
 }
 
+TEST(RawCodegen, Server_InvokeUnaryRpc) {
+  PW_RAW_TEST_METHOD_CONTEXT(test::TestService, TestRpc) context;
+
+  std::byte buffer[64];
+  protobuf::NestedEncoder encoder(buffer);
+  test::TestRequest::Encoder test_request(&encoder);
+  test_request.WriteInteger(123);
+  test_request.WriteStatusCode(Status::Ok());
+
+  auto sws = context.call(encoder.Encode().value());
+  EXPECT_EQ(Status::Ok(), sws.status());
+
+  protobuf::Decoder decoder(context.response());
+
+  while (decoder.Next().ok()) {
+    switch (static_cast<test::TestResponse::Fields>(decoder.FieldNumber())) {
+      case test::TestResponse::Fields::VALUE: {
+        int32_t value;
+        decoder.ReadInt32(&value);
+        EXPECT_EQ(value, 124);
+        break;
+      }
+    }
+  }
+}
+
+TEST(RawCodegen, Server_InvokeServerStreamingRpc) {
+  PW_RAW_TEST_METHOD_CONTEXT(test::TestService, TestStreamRpc) context;
+
+  std::byte buffer[64];
+  protobuf::NestedEncoder encoder(buffer);
+  test::TestRequest::Encoder test_request(&encoder);
+  test_request.WriteInteger(5);
+  test_request.WriteStatusCode(Status::Unauthenticated());
+
+  context.call(encoder.Encode().value());
+  EXPECT_TRUE(context.done());
+  EXPECT_EQ(Status::Unauthenticated(), context.status());
+  EXPECT_EQ(context.total_responses(), 5u);
+
+  protobuf::Decoder decoder(context.responses().back());
+  while (decoder.Next().ok()) {
+    switch (
+        static_cast<test::TestStreamResponse::Fields>(decoder.FieldNumber())) {
+      case test::TestStreamResponse::Fields::NUMBER: {
+        int32_t value;
+        decoder.ReadInt32(&value);
+        EXPECT_EQ(value, 4);
+        break;
+      }
+      case test::TestStreamResponse::Fields::CHUNK:
+        FAIL();
+        break;
+    }
+  }
+}
+
 }  // namespace
 }  // namespace pw::rpc
diff --git a/pw_rpc/raw/public/pw_rpc/internal/raw_service_method_traits.h b/pw_rpc/raw/public/pw_rpc/internal/raw_service_method_traits.h
new file mode 100644
index 0000000..e304fc2
--- /dev/null
+++ b/pw_rpc/raw/public/pw_rpc/internal/raw_service_method_traits.h
@@ -0,0 +1,27 @@
+// Copyright 2020 The Pigweed Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not
+// use this file except in compliance with the License. You may obtain a copy of
+// the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations under
+// the License.
+#pragma once
+
+#include "pw_rpc/internal/raw_method_union.h"
+#include "pw_rpc/internal/service_method_traits.h"
+
+namespace pw::rpc::internal {
+
+template <auto impl_method, uint32_t method_id>
+using RawServiceMethodTraits =
+    ServiceMethodTraits<&MethodBaseService<impl_method>::RawMethodFor,
+                        impl_method,
+                        method_id>;
+
+}  // namespace pw::rpc::internal
diff --git a/pw_rpc/raw/public/pw_rpc/raw_test_method_context.h b/pw_rpc/raw/public/pw_rpc/raw_test_method_context.h
new file mode 100644
index 0000000..732a2b2
--- /dev/null
+++ b/pw_rpc/raw/public/pw_rpc/raw_test_method_context.h
@@ -0,0 +1,312 @@
+// Copyright 2020 The Pigweed Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not
+// use this file except in compliance with the License. You may obtain a copy of
+// the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations under
+// the License.
+#pragma once
+
+#include <type_traits>
+
+#include "pw_assert/light.h"
+#include "pw_bytes/span.h"
+#include "pw_containers/vector.h"
+#include "pw_rpc/channel.h"
+#include "pw_rpc/internal/hash.h"
+#include "pw_rpc/internal/packet.h"
+#include "pw_rpc/internal/raw_service_method_traits.h"
+#include "pw_rpc/internal/server.h"
+
+namespace pw::rpc {
+
+// Declares a context object that may be used to invoke an RPC. The context is
+// 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_RAW_TEST_METHOD_CONTEXT(my::CoolService, TheMethod) context;
+//   EXPECT_EQ(Status::Ok(), context.call(encoded_request).status());
+//   EXPECT_EQ(0,
+//             std::memcmp(encoded_response,
+//                         context.response().data(),
+//                         sizeof(encoded_response)));
+//
+// For a server streaming RPC, context.call(request) invokes the method. As in a
+// normal RPC, the method completes when the ServerWriter's Finish method is
+// called (or it goes out of scope).
+//
+//   PW_RAW_TEST_METHOD_CONTEXT(my::CoolService, TheStreamingMethod) context;
+//   context.call(encoded_response);
+//
+//   EXPECT_TRUE(context.done());  // Check that the RPC completed
+//   EXPECT_EQ(Status::Ok(), context.status());  // Check the status
+//
+//   EXPECT_EQ(3u, context.responses().size());
+//   ByteSpan& response = context.responses()[0];  // check individual responses
+//
+//   for (ByteSpan& response : context.responses()) {
+//     // iterate over the responses
+//   }
+//
+// PW_RAW_TEST_METHOD_CONTEXT forwards its constructor arguments to the
+// underlying service. For example:
+//
+//   PW_RAW_TEST_METHOD_CONTEXT(MyService, Go) context(service, args);
+//
+// PW_RAW_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_RAW_TEST_METHOD_CONTEXT(MyService, BestMethod, 3, 256) context;
+//   ASSERT_EQ(3u, context.responses().max_size());
+//
+#define PW_RAW_TEST_METHOD_CONTEXT(service, method, ...)              \
+  ::pw::rpc::RawTestMethodContext<&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 RawTestMethodContext;
+
+// Internal classes that implement RawTestMethodContext.
+namespace internal::test::raw {
+
+// A ChannelOutput implementation that stores the outgoing payloads and status.
+template <size_t output_size>
+class MessageOutput final : public ChannelOutput {
+ public:
+  using ResponseBuffer = std::array<std::byte, output_size>;
+
+  MessageOutput(Vector<ByteSpan>& responses,
+                Vector<ResponseBuffer>& buffers,
+                ByteSpan packet_buffer)
+      : ChannelOutput("internal::test::raw::MessageOutput"),
+        responses_(responses),
+        buffers_(buffers),
+        packet_buffer_(packet_buffer) {
+    clear();
+  }
+
+  Status last_status() const { return last_status_; }
+  void set_last_status(Status status) { last_status_ = status; }
+
+  size_t total_responses() const { return total_responses_; }
+
+  bool stream_ended() const { return stream_ended_; }
+
+  void clear() {
+    responses_.clear();
+    buffers_.clear();
+    total_responses_ = 0;
+    stream_ended_ = false;
+    last_status_ = Status::Unknown();
+  }
+
+ private:
+  ByteSpan AcquireBuffer() override { return packet_buffer_; }
+
+  Status SendAndReleaseBuffer(size_t size) override;
+
+  Vector<ByteSpan>& responses_;
+  Vector<ResponseBuffer>& buffers_;
+  ByteSpan packet_buffer_;
+  size_t total_responses_;
+  bool stream_ended_;
+  Status last_status_;
+};
+
+// Collects everything needed to invoke a particular RPC.
+template <auto method,
+          uint32_t method_id,
+          size_t max_responses,
+          size_t output_size>
+struct InvocationContext {
+  template <typename... Args>
+  InvocationContext(Args&&... args)
+      : output(responses, buffers, packet_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,
+             RawServiceMethodTraits<method, method_id>::method()) {}
+
+  using ResponseBuffer = std::array<std::byte, output_size>;
+  using Service = typename RawServiceMethodTraits<method, method_id>::Service;
+
+  MessageOutput<output_size> output;
+  rpc::Channel channel;
+  rpc::Server server;
+  Service service;
+  Vector<ByteSpan, max_responses> responses;
+  Vector<ResponseBuffer, max_responses> buffers;
+  std::array<std::byte, output_size> packet_buffer = {};
+  internal::ServerCall call;
+};
+
+// Method invocation context for a unary RPC. Returns the status in call() and
+// provides the response through the response() method.
+template <auto method, uint32_t method_id, size_t output_size>
+class UnaryContext {
+ private:
+  using Context = InvocationContext<method, method_id, 1, output_size>;
+  Context ctx_;
+
+ public:
+  template <typename... Args>
+  UnaryContext(Args&&... args) : ctx_(std::forward<Args>(args)...) {}
+
+  typename Context::Service& service() { return ctx_.service; }
+
+  // Invokes the RPC with the provided request. Returns RPC's StatusWithSize.
+  StatusWithSize call(ConstByteSpan request) {
+    ctx_.output.clear();
+    ctx_.buffers.emplace_back();
+    ctx_.buffers.back() = {};
+    ctx_.responses.emplace_back();
+    auto& response = ctx_.responses.back();
+    response = {ctx_.buffers.back().data(), ctx_.buffers.back().size()};
+    auto sws = (ctx_.service.*method)(ctx_.call.context(), request, response);
+    response = response.first(sws.size());
+    return sws;
+  }
+
+  // Gives access to the RPC's response.
+  ConstByteSpan response() const {
+    PW_ASSERT(ctx_.responses.size() > 0u);
+    return ctx_.responses.back();
+  }
+};
+
+// Method invocation context for a server streaming RPC.
+template <auto method,
+          uint32_t method_id,
+          size_t max_responses,
+          size_t output_size>
+class ServerStreamingContext {
+ private:
+  using Context =
+      InvocationContext<method, method_id, max_responses, output_size>;
+  Context ctx_;
+
+ public:
+  template <typename... Args>
+  ServerStreamingContext(Args&&... args) : ctx_(std::forward<Args>(args)...) {}
+
+  typename Context::Service& service() { return ctx_.service; }
+
+  // Invokes the RPC with the provided request.
+  void call(ConstByteSpan request) {
+    ctx_.output.clear();
+    BaseServerWriter server_writer(ctx_.call);
+    return (ctx_.service.*method)(ctx_.call.context(),
+                                  request,
+                                  static_cast<RawServerWriter&>(server_writer));
+  }
+
+  // Returns a server writer which writes responses into the context's buffer.
+  // This should not be called alongside call(); use one or the other.
+  RawServerWriter writer() {
+    ctx_.output.clear();
+    BaseServerWriter server_writer(ctx_.call);
+    return std::move(static_cast<RawServerWriter&>(server_writer));
+  }
+
+  // Returns the responses that have been recorded. The maximum number of
+  // responses is responses().max_size(). responses().back() is always the most
+  // recent response, even if total_responses() > responses().max_size().
+  const Vector<ByteSpan>& responses() const { return ctx_.responses; }
+
+  // The total number of responses sent, which may be larger than
+  // responses.max_size().
+  size_t total_responses() const { return ctx_.output.total_responses(); }
+
+  // True if the stream has terminated.
+  bool done() const { return ctx_.output.stream_ended(); }
+
+  // The status of the stream. Only valid if done() is true.
+  Status status() const {
+    PW_ASSERT(done());
+    return ctx_.output.last_status();
+  }
+};
+
+// Alias to select the type of the context object to use based on which type of
+// RPC it is for.
+template <auto method, uint32_t method_id, size_t responses, size_t output_size>
+using Context = std::tuple_element_t<
+    static_cast<size_t>(MethodTraits<decltype(method)>::kType),
+    std::tuple<UnaryContext<method, method_id, output_size>,
+               ServerStreamingContext<method, method_id, responses, output_size>
+               // TODO(hepler): Support client and bidi streaming
+               >>;
+
+template <size_t output_size>
+Status MessageOutput<output_size>::SendAndReleaseBuffer(size_t size) {
+  PW_ASSERT(!stream_ended_);
+
+  if (size == 0u) {
+    return Status::Ok();
+  }
+
+  Result<internal::Packet> result =
+      internal::Packet::FromBuffer(std::span(packet_buffer_.data(), size));
+  PW_ASSERT(result.ok());
+
+  last_status_ = result.value().status();
+
+  switch (result.value().type()) {
+    case internal::PacketType::RESPONSE: {
+      // If we run out of space, the back message is always the most recent.
+      buffers_.emplace_back();
+      buffers_.back() = {};
+      auto response = result.value().payload();
+      std::memcpy(&buffers_.back(), response.data(), response.size());
+      responses_.emplace_back();
+      responses_.back() = {buffers_.back().data(), response.size()};
+      total_responses_ += 1;
+      break;
+    }
+    case internal::PacketType::SERVER_STREAM_END:
+      stream_ended_ = true;
+      break;
+    default:
+      PW_CRASH("Unhandled PacketType");
+  }
+  return Status::Ok();
+}
+
+}  // namespace internal::test::raw
+
+template <auto method,
+          uint32_t method_id,
+          size_t max_responses,
+          size_t output_size_bytes>
+class RawTestMethodContext
+    : public internal::test::raw::
+          Context<method, method_id, max_responses, output_size_bytes> {
+ public:
+  // Forwards constructor arguments to the service class.
+  template <typename... ServiceArgs>
+  RawTestMethodContext(ServiceArgs&&... service_args)
+      : internal::test::raw::
+            Context<method, method_id, max_responses, output_size_bytes>(
+                std::forward<ServiceArgs>(service_args)...) {}
+};
+
+}  // namespace pw::rpc