pw_rpc2/cpp/client_test.cc

// Copyright 2026 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.

#include "pw_rpc2/client.h"

#include <cstddef>
#include <optional>
#include <utility>
#include <variant>

#include "echo_pb/echo.rpc.pwpb.h"
#include "pw_allocator/testing.h"
#include "pw_assert/check.h"
#include "pw_async2/dispatcher_for_test.h"
#include "pw_async2/try.h"
#include "pw_async2/value_future.h"
#include "pw_rpc2/internal/handshake.h"
#include "pw_rpc2/internal/packet.h"
#include "pw_rpc2/internal/test_utils.h"
#include "pw_thread/test_thread_context.h"
#include "pw_thread/thread.h"
#include "pw_transport/transport.h"
#include "pw_unit_test/framework.h"

namespace pw::rpc2 {
namespace {

class MockConnector : public transport::ReliableDatagramConnector {
 public:
  MockConnector() = default;

  transport::ReliableDatagramConnector::ConnectFuture Connect() override {
    return connect_provider_.Get();
  }

  void ResolveConnect(Result<transport::ReliableDatagramSocket> res) {
    connect_provider_.Resolve(std::move(res));
  }

 private:
  async2::ValueProvider<Result<transport::ReliableDatagramSocket>>
      connect_provider_;
};

class ClientFutureTestTask : public async2::Task {
 public:
  explicit ClientFutureTestTask(ClientFuture future)
      : fut_(std::move(future)) {}

  async2::Poll<> DoPend(async2::Context& cx) override {
    auto poll = fut_.Pend(cx);
    if (poll.IsPending()) {
      return async2::Pending();
    }
    result_ = std::move(*poll);
    return async2::Ready();
  }

  std::optional<Result<Client>>& result() { return result_; }
  const std::optional<Result<Client>>& result() const { return result_; }

  ~ClientFutureTestTask() override {
    Deregister();
    // A `Client` must be closed before it is destroyed. If nothing took
    // this one, close it here.
    if (result_.has_value() && result_->ok()) {
      auto closed = (*result_)->Close();
    }
  }

 private:
  ClientFuture fut_;
  std::optional<Result<Client>> result_;
};

// Awaits a `ControlFuture` on the dispatcher. A future only makes progress
// while a task is pending it, so tests cannot simply run the dispatcher and
// inspect the future.
class ControlTask : public async2::Task {
 public:
  explicit ControlTask(ControlFuture&& future)
      : async2::Task(PW_ASYNC_TASK_NAME("ControlTask")),
        future_(std::move(future)) {}

  async2::Poll<> DoPend(async2::Context& cx) override {
    PW_TRY_READY_ASSIGN(Status status, future_.Pend(cx));
    status_ = status;
    return async2::Ready();
  }

  bool done() const { return status_.has_value(); }
  Status status() const { return status_.value_or(Status::Unknown()); }

 private:
  ControlFuture future_;
  std::optional<Status> status_;
};

// Pumps the handshake against a mock connection and returns the connected
// client. `task` must outlive the returned client's creation, but is
// deregistered before returning.
Client ConnectClient(async2::DispatcherForTest& dispatcher,
                     Allocator& allocator,
                     transport::ReliableDatagramSocket connection,
                     test::MockConnection& raw_connection,
                     std::optional<ClientFutureTestTask>& task) {
  task.emplace(Client::Create(dispatcher, allocator, std::move(connection)));
  dispatcher.Post(*task);
  dispatcher.RunUntilStalled();

  Buf response =
      Buf::Allocate(allocator, internal::HandshakePacket::kWireSizeBytes);
  PW_CHECK(!response.empty());
  auto encoded =
      internal::HandshakePacket(internal::HandshakePacket::Type::kSynAck)
          .Encode(std::move(response));
  PW_CHECK(encoded.ok());
  raw_connection.SetNextRead(std::move(*encoded));

  dispatcher.RunUntilStalled();
  PW_CHECK(task->result().has_value());
  PW_CHECK(task->result()->ok());
  Client client = std::move(task->result()->value());
  task->Deregister();
  return client;
}

TEST(ClientFutureTest, CreateFromTransportSequentialConnectAndHandshake) {
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  MockConnector transport;
  auto [conn, raw_conn] = test::MakeMockConnection(allocator);

  auto fut = Client::Create(dispatcher, allocator, transport);
  ClientFutureTestTask task(std::move(fut));
  dispatcher.Post(task);

  // Poll 1: Awaiting transport connection.
  dispatcher.RunUntilStalled();
  EXPECT_FALSE(task.result().has_value());

  // Resolve transport connection. Now ClientFuture should move to handshake
  // stage and send HandshakePacket.
  transport.ResolveConnect(conn);
  dispatcher.RunUntilStalled();
  EXPECT_FALSE(task.result().has_value());
  EXPECT_EQ(raw_conn->commit_count(), 1u);

  // Provide server handshake response.
  Buf resp =
      Buf::Allocate(allocator, internal::HandshakePacket::kWireSizeBytes);
  ASSERT_FALSE(resp.empty());
  auto hs_pkt =
      internal::HandshakePacket(internal::HandshakePacket::Type::kSynAck);
  auto enc_res = hs_pkt.Encode(std::move(resp));
  ASSERT_TRUE(enc_res.ok());
  raw_conn->SetNextRead(std::move(enc_res.value()));

  // Poll 2: ClientFuture completes handshake (sends ACK) and yields handshaked
  // Client.
  dispatcher.RunUntilStalled();
  ASSERT_TRUE(task.result().has_value());
  ASSERT_TRUE(task.result()->ok());
  EXPECT_EQ(raw_conn->commit_count(), 2u);

  Client client = std::move(task.result()->value());
  EXPECT_TRUE(client.is_open());

  // Check that the client can create calls after handshake completes
  auto call = internal::CallAccess::CreateCall(client);
  EXPECT_NE(call, nullptr);
  call = nullptr;

  test::CloseClient(client, dispatcher);
  task.Deregister();
}

TEST(ClientFutureTest, CreateFromConnectionDirectHandshake) {
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  auto [conn, raw_conn] = test::MakeMockConnection(allocator);

  auto fut = Client::Create(dispatcher, allocator, conn);
  ClientFutureTestTask task(std::move(fut));
  dispatcher.Post(task);

  // Poll 1: Awaiting handshake response.
  dispatcher.RunUntilStalled();
  EXPECT_FALSE(task.result().has_value());
  EXPECT_EQ(raw_conn->commit_count(), 1u);

  // Provide server handshake response.
  Buf resp =
      Buf::Allocate(allocator, internal::HandshakePacket::kWireSizeBytes);
  ASSERT_FALSE(resp.empty());
  auto hs_pkt =
      internal::HandshakePacket(internal::HandshakePacket::Type::kSynAck);
  auto enc_res = hs_pkt.Encode(std::move(resp));
  ASSERT_TRUE(enc_res.ok());
  raw_conn->SetNextRead(std::move(enc_res.value()));

  // Poll 2: Completes handshake and yields ready Client.
  dispatcher.RunUntilStalled();
  ASSERT_TRUE(task.result().has_value());
  ASSERT_TRUE(task.result()->ok());
  EXPECT_EQ(raw_conn->commit_count(), 2u);

  Client client = std::move(task.result()->value());
  EXPECT_TRUE(client.is_open());

  test::CloseClient(client, dispatcher);
  task.Deregister();
}

TEST(ClientFutureTest, CreateFromTransportConnectFailure) {
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  MockConnector transport;
  auto fut = Client::Create(dispatcher, allocator, transport);
  ClientFutureTestTask task(std::move(fut));
  dispatcher.Post(task);

  dispatcher.RunUntilStalled();
  EXPECT_FALSE(task.result().has_value());

  // Fail transport connection.
  transport.ResolveConnect(Status::Unavailable());
  dispatcher.RunUntilStalled();

  ASSERT_TRUE(task.result().has_value());
  EXPECT_FALSE(task.result()->ok());
  EXPECT_EQ(task.result()->status(), Status::Unavailable());

  task.Deregister();
}

TEST(ClientFutureTest, CreateFromConnectionHandshakeFailure) {
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  auto [conn, raw_conn] = test::MakeMockConnection(allocator);

  auto fut = Client::Create(dispatcher, allocator, conn);
  ClientFutureTestTask task(std::move(fut));
  dispatcher.Post(task);

  dispatcher.RunUntilStalled();
  EXPECT_FALSE(task.result().has_value());

  // Provide corrupt handshake packet response.
  Buf resp = Buf::Allocate(allocator, 4);
  ASSERT_FALSE(resp.empty());
  raw_conn->SetNextRead(std::move(resp));

  dispatcher.RunUntilStalled();
  ASSERT_TRUE(task.result().has_value());
  EXPECT_FALSE(task.result()->ok());
  EXPECT_EQ(task.result()->status(), Status::DataLoss());

  task.Deregister();
}

TEST(ClientTest, CallIdAllocationAndCopyMoveSemantics) {
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  auto [conn, raw_conn] = test::MakeMockConnection(allocator);

  auto fut = Client::Create(dispatcher, allocator, conn);
  ClientFutureTestTask task(std::move(fut));
  dispatcher.Post(task);

  dispatcher.RunUntilStalled();
  EXPECT_FALSE(task.result().has_value());

  pw::Buf resp =
      pw::Buf::Allocate(allocator, internal::HandshakePacket::kWireSizeBytes);
  auto hs_pkt =
      internal::HandshakePacket(internal::HandshakePacket::Type::kSynAck);
  auto enc_res = hs_pkt.Encode(std::move(resp));
  ASSERT_TRUE(enc_res.ok());
  raw_conn->SetNextRead(std::move(enc_res.value()));

  dispatcher.RunUntilStalled();
  ASSERT_TRUE(task.result().has_value() && task.result()->ok());
  EXPECT_EQ(raw_conn->commit_count(), 2u);

  Client client = std::move(task.result()->value());

  // Moving empties the source.
  Client moved_client(std::move(client));
  EXPECT_TRUE(moved_client.is_open());
  EXPECT_FALSE(client.is_open());  // NOLINT(bugprone-use-after-move)

  // Copying keeps both clients open and referring to the same connection.
  Client copied_client(moved_client);
  EXPECT_TRUE(copied_client.is_open());
  EXPECT_TRUE(moved_client.is_open());

  // ServiceClients store a Client and expose it via client(), and every copy
  // draws call IDs from the one underlying connection. A ServiceClient is only
  // constructed through a derived service client.
  class TestServiceClient : public ServiceClient {
   public:
    TestServiceClient(const Client& client, uint32_t service_id)
        : ServiceClient(client, service_id) {}
  };
  TestServiceClient ref1(copied_client, 1);
  TestServiceClient ref2 = ref1;
  EXPECT_TRUE(ref1.is_open());
  EXPECT_TRUE(ref1.client().is_open());
  auto call1 = internal::CallAccess::CreateCall(ref1.client());
  auto call2 = internal::CallAccess::CreateCall(ref2.client());
  auto call3 = internal::CallAccess::CreateCall(copied_client);
  ASSERT_NE(call1, nullptr);
  ASSERT_NE(call2, nullptr);
  ASSERT_NE(call3, nullptr);
  EXPECT_NE(call1->call_id(), call2->call_id());
  EXPECT_NE(call2->call_id(), call3->call_id());

  // Drop the calls before closing the client.
  call1 = nullptr;
  call2 = nullptr;
  call3 = nullptr;

  test::CloseClient(ref1.client(), dispatcher);
  EXPECT_FALSE(moved_client.is_open());
  EXPECT_FALSE(copied_client.is_open());
  EXPECT_FALSE(ref2.is_open());
  task.Deregister();
}

TEST(MockPeerTest, SkipsHandshakeEntirely) {
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  auto peer = test::MakeMockPeer(dispatcher, allocator);

  // Client::Create() would have written a kSyn here and blocked awaiting a
  // kSynAck (see CreateFromConnectionDirectHandshake above, which expects
  // commit_count 1 then 2). MakeMockPeer writes nothing at all, so the
  // first recorded packet is whatever the first real call emits.
  dispatcher.RunUntilStalled();
  EXPECT_EQ(peer.connection().commit_count(), 0u);
  EXPECT_EQ(peer.connection().written_packet_count(), 0u);
}

TEST(MockPeerTest, IsImmediatelyUsable) {
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  auto peer = test::MakeMockPeer(dispatcher, allocator);
  const Client& client = peer.client();

  // No handshake had to be pumped for the client to allocate calls.
  EXPECT_TRUE(client.is_open());
  auto call = internal::CallAccess::CreateCall(client);
  EXPECT_NE(call, nullptr);
}

TEST(MockPeerTest, ObservesNoCallsBeforeAnyCall) {
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  auto peer = test::MakeMockPeer(dispatcher, allocator);

  EXPECT_TRUE(peer.pending_calls().empty());
}

// The typed payload path is backend-agnostic: `MockPeer` only ever calls the
// generic `Size`/`Serialize`/`Deserialize` dispatchers, and the pwpb
// `SerializerFor` specialization arrives via the generated header included
// above. These tests pin that down against a real generated client.

using EchoRequest = echo::pwpb::EchoRequest::Message;
using EchoResponse = echo::pwpb::EchoResponse::Message;

class EchoUnaryTask : public async2::Task {
 public:
  using Future =
      echo::pwpb::EchoService::UnaryCallFuture<EchoRequest, EchoResponse>;

  explicit EchoUnaryTask(Future fut) : fut_(std::move(fut)) {}

  async2::Poll<> DoPend(async2::Context& cx) override {
    auto poll = fut_.Pend(cx);
    if (poll.IsPending()) {
      return async2::Pending();
    }
    result_ = std::move(*poll);
    return async2::Ready();
  }

  const std::optional<Result<EchoResponse>>& result() const { return result_; }

 private:
  Future fut_;
  std::optional<Result<EchoResponse>> result_;
};

TEST(MockPeerTest, DecodesTypedRequestAndSendsTypedResponse) {
  // Allocator first: the dispatcher destroys still-posted tasks living in it.
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  auto peer = test::MakeMockPeer(dispatcher, allocator);
  echo::pwpb::EchoService::Client client(peer.client());

  EchoUnaryTask task(client.EchoUnary(EchoRequest{.val = 42}));
  dispatcher.Post(task);
  dispatcher.RunUntilStalled();

  // Read side: the recorded request payload decodes back to the message.
  auto call = peer.ExpectInvocation<echo::pwpb::EchoService::EchoUnary>();
  auto decoded = call.request();
  ASSERT_TRUE(decoded.ok());
  EXPECT_EQ(decoded->val, 42u);

  // The call is still open until the peer answers it.
  EXPECT_FALSE(task.result().has_value());

  // Send side: a typed response is serialized and delivered to the client.
  call.Finish(EchoResponse{.val = 99});

  ASSERT_TRUE(task.result().has_value());
  ASSERT_TRUE(task.result()->ok());
  EXPECT_EQ(task.result()->value().val, 99u);

  task.Deregister();
}

TEST(MockPeerTest, ClaimsEachInvocationExactlyOnce) {
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  auto peer = test::MakeMockPeer(dispatcher, allocator);
  echo::pwpb::EchoService::Client client(peer.client());

  EchoUnaryTask first(client.EchoUnary(EchoRequest{.val = 1}));
  dispatcher.Post(first);
  peer.ExpectInvocation<echo::pwpb::EchoService::EchoUnary>().Finish(
      EchoResponse{.val = 11});
  ASSERT_TRUE(first.result().has_value());
  EXPECT_EQ(first.result()->value().val, 11u);

  // The first call is claimed, so the second ExpectInvocation() sees only the
  // new one rather than tripping on a stale request.
  EchoUnaryTask second(client.EchoUnary(EchoRequest{.val = 2}));
  dispatcher.Post(second);
  peer.ExpectInvocation<echo::pwpb::EchoService::EchoUnary>().Finish(
      EchoResponse{.val = 22});
  ASSERT_TRUE(second.result().has_value());
  EXPECT_EQ(second.result()->value().val, 22u);

  EXPECT_TRUE(peer.pending_calls().empty());

  first.Deregister();
  second.Deregister();
}

TEST(MockPeerTest, TypedDecodeFailsOnGarbagePayload) {
  // Decoding a valid payload as the wrong shape must report an error rather
  // than silently yielding a default-constructed message.
  const std::byte garbage[] = {std::byte{0xFF}, std::byte{0xFF}};
  EXPECT_FALSE(test::PayloadAs<EchoRequest>(garbage).ok());
}

// Closing a client is the only way to release its connection.
TEST(ClientTest, CloseMarksClientClosedAndResolves) {
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  auto [conn, raw_conn] = test::MakeMockConnection(allocator);
  std::optional<ClientFutureTestTask> connect_task;
  Client client =
      ConnectClient(dispatcher, allocator, conn, *raw_conn, connect_task);
  ASSERT_TRUE(client.is_open());

  // `Close()` hands the close request off to the dispatcher immediately, so
  // `is_open()` becomes false as soon as `Close()` returns.
  ControlTask close_task(client.Close());
  EXPECT_FALSE(client.is_open());

  dispatcher.Post(close_task);
  dispatcher.RunUntilStalled();
  EXPECT_TRUE(close_task.done());
  EXPECT_EQ(close_task.status(), OkStatus());
}

// Closing an empty client is a no-op rather than an error, so moved-from and
// already-closed clients can be closed again unconditionally.
TEST(ClientTest, CloseOnAnEmptyClientResolvesImmediately) {
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  auto [conn, raw_conn] = test::MakeMockConnection(allocator);
  std::optional<ClientFutureTestTask> connect_task;
  Client client =
      ConnectClient(dispatcher, allocator, conn, *raw_conn, connect_task);

  test::CloseClient(client, dispatcher);
  ASSERT_FALSE(client.is_open());

  ControlTask close_task(client.Close());
  dispatcher.Post(close_task);
  dispatcher.RunUntilStalled();
  EXPECT_TRUE(close_task.done());
  EXPECT_EQ(close_task.status(), OkStatus());
}

// Closing tears down the connection, which aborts every call riding on it.
TEST(ClientTest, CloseCancelsInFlightCalls) {
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  auto [conn, raw_conn] = test::MakeMockConnection(allocator);
  std::optional<ClientFutureTestTask> connect_task;
  Client client =
      ConnectClient(dispatcher, allocator, conn, *raw_conn, connect_task);

  echo::pwpb::EchoService::Client stub{client};
  EchoUnaryTask call(stub.EchoUnary(EchoRequest{.val = 7}));
  dispatcher.Post(call);
  dispatcher.RunUntilStalled();
  ASSERT_FALSE(call.result().has_value());

  test::CloseClient(client, dispatcher);

  ASSERT_TRUE(call.result().has_value());
  EXPECT_EQ(call.result()->status(), Status::Cancelled());

  call.Deregister();
}

// `Close()` and `CloseBlocking()` are the only client operations that may run
// off the dispatcher thread; the blocking form is for threads that are not
// driving the dispatcher themselves.
TEST(ClientTest, BlockingCloseRunsFromAnotherThread) {
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  auto [conn, raw_conn] = test::MakeMockConnection(allocator);
  std::optional<ClientFutureTestTask> connect_task;
  Client client =
      ConnectClient(dispatcher, allocator, conn, *raw_conn, connect_task);

  pw::thread::test::TestThreadContext context;
  pw::Thread thread(context.options(), [&] { client.CloseBlocking(); });

  dispatcher.AllowBlocking();
  dispatcher.RunToCompletion();
  thread.join();

  EXPECT_FALSE(client.is_open());
}

TEST(ClientTest, DroppingOneOfSeveralHandlesKeepsConnectionOpen) {
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  auto [conn, raw_conn] = test::MakeMockConnection(allocator);
  std::optional<ClientFutureTestTask> connect_task;
  std::optional<Client> client =
      ConnectClient(dispatcher, allocator, conn, *raw_conn, connect_task);

  echo::pwpb::EchoService::Client stub{*client};
  EchoUnaryTask call(stub.EchoUnary(EchoRequest{.val = 7}));
  dispatcher.Post(call);
  dispatcher.RunUntilStalled();
  ASSERT_FALSE(call.result().has_value());

  // Destroying one Client keeps the connection open for the remaining handles
  // (such as `stub`).
  client.reset();
  dispatcher.RunUntilStalled();

  EXPECT_TRUE(stub.is_open());
  EXPECT_FALSE(raw_conn->is_closed());
  EXPECT_FALSE(call.result().has_value());

  test::CloseClient(stub.client(), dispatcher);
  ASSERT_TRUE(call.result().has_value());
  EXPECT_EQ(call.result()->status(), Status::Cancelled());

  call.Deregister();
}

// Once no handle can reach the connection, nothing could ever close it, so
// dropping the last handle closes it.
TEST(ClientTest, DroppingLastHandleClosesConnection) {
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  auto [conn, raw_conn] = test::MakeMockConnection(allocator);
  std::optional<ClientFutureTestTask> connect_task;
  std::optional<Client> client =
      ConnectClient(dispatcher, allocator, conn, *raw_conn, connect_task);
  std::optional<echo::pwpb::EchoService::Client> stub(std::in_place, *client);

  EchoUnaryTask call(stub->EchoUnary(EchoRequest{.val = 7}));
  dispatcher.Post(call);
  dispatcher.RunUntilStalled();
  ASSERT_FALSE(call.result().has_value());

  client.reset();
  dispatcher.RunUntilStalled();
  EXPECT_FALSE(raw_conn->is_closed());
  EXPECT_FALSE(call.result().has_value());

  // `stub` holds the last handle.
  stub.reset();
  dispatcher.RunUntilStalled();

  EXPECT_TRUE(raw_conn->is_closed());
  ASSERT_TRUE(call.result().has_value());
  EXPECT_EQ(call.result()->status(), Status::Cancelled());

  call.Deregister();
}

TEST(ClientFutureTest, DestroyedUnpendedFromConnectionClosesConnection) {
  allocator::test::AllocatorForTest<16384> allocator;
  async2::DispatcherForTest dispatcher;

  auto [conn, raw_conn] = test::MakeMockConnection(allocator);
  {
    auto fut = Client::Create(dispatcher, allocator, conn);
    EXPECT_TRUE(fut.is_pendable());
  }
  EXPECT_TRUE(raw_conn->is_closed());
}

}  // namespace
}  // namespace pw::rpc2