From e2115152c899eb8cd7f79915105de429724721dc Mon Sep 17 00:00:00 2001
From: Andrew Noyes <andrew@weaselab.dev>
Date: Mon, 15 Sep 2025 15:48:10 -0400
Subject: [PATCH] Add tests for shutdown vs close

---
 tests/test_server.cpp | 112 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 112 insertions(+)

diff --git a/tests/test_server.cpp b/tests/test_server.cpp
index 35739d4..e5d9dce 100644
--- a/tests/test_server.cpp
+++ b/tests/test_server.cpp
@@ -55,3 +55,115 @@ TEST_CASE("Echo test") {
   server->shutdown();
   runThread.join();
 }
+
+struct ShutdownTestHandler : ConnectionHandler {
+  Arena arena;
+  std::span<std::string_view> reply;
+  WeakRef<MessageSender> wconn;
+  std::latch received_data{1};
+  ConnectionShutdown shutdown_mode = ConnectionShutdown::None;
+  std::atomic<bool> connection_closed{false};
+
+  void on_data_arrived(std::string_view data, Connection &conn) override {
+    reply = arena.allocate_span<std::string_view>(1);
+    reply[0] = arena.copy_string(data);
+    wconn = conn.get_weak_ref();
+    received_data.count_down();
+  }
+
+  void on_connection_closed(Connection &) override { connection_closed = true; }
+};
+
+TEST_CASE("Connection shutdown write-only mode") {
+  ShutdownTestHandler handler;
+  handler.shutdown_mode = ConnectionShutdown::WriteOnly;
+  weaseldb::Config config;
+  auto server = Server::create(config, handler, {});
+  int fd = server->create_local_connection();
+
+  auto runThread = std::thread{[&]() { server->run(); }};
+
+  // Send data to trigger handler
+  int w = write(fd, "test", 4);
+  REQUIRE(w == 4);
+
+  handler.received_data.wait();
+
+  // Send response with write shutdown
+  if (auto conn = handler.wconn.lock()) {
+    auto *conn_ptr = static_cast<Connection *>(conn.get());
+    conn_ptr->append_bytes(std::exchange(handler.reply, {}),
+                           std::move(handler.arena),
+                           ConnectionShutdown::WriteOnly);
+  } else {
+    REQUIRE(false);
+  }
+
+  // Read the response
+  char buf[5];
+  buf[4] = 0;
+  int r = read(fd, buf, 4);
+  REQUIRE(r == 4);
+  CHECK(std::string(buf) == "test");
+
+  // After write shutdown, we should get EOF when trying to read more
+  char extra_buf[1];
+  int eof_result = read(fd, extra_buf, 1);
+  CHECK(eof_result == 0); // EOF indicates successful write shutdown
+
+  // Connection should still be alive (not closed) after write shutdown
+  // We can verify this by checking that we can still write to the socket
+  int write_result = write(fd, "x", 1);
+  CHECK(write_result == 1); // Should succeed - connection still alive
+  CHECK(handler.connection_closed.load() ==
+        false); // Connection should still be alive
+
+  close(fd);
+  server->shutdown();
+  runThread.join();
+}
+
+TEST_CASE("Connection shutdown full mode") {
+  ShutdownTestHandler handler;
+  handler.shutdown_mode = ConnectionShutdown::Full;
+  weaseldb::Config config;
+  auto server = Server::create(config, handler, {});
+  int fd = server->create_local_connection();
+
+  auto runThread = std::thread{[&]() { server->run(); }};
+
+  // Send data to trigger handler
+  int w = write(fd, "test", 4);
+  REQUIRE(w == 4);
+
+  handler.received_data.wait();
+
+  // Send response with full shutdown
+  if (auto conn = handler.wconn.lock()) {
+    auto *conn_ptr = static_cast<Connection *>(conn.get());
+    conn_ptr->append_bytes(std::exchange(handler.reply, {}),
+                           std::move(handler.arena), ConnectionShutdown::Full);
+  } else {
+    REQUIRE(false);
+  }
+
+  // Read the response - connection should close after this
+  char buf[5];
+  buf[4] = 0;
+  int r = read(fd, buf, 4);
+  REQUIRE(r == 4);
+  CHECK(std::string(buf) == "test");
+
+  // Connection should be closed by server (full shutdown)
+  char extra_buf[1];
+  int close_result = read(fd, extra_buf, 1);
+  CHECK(close_result == 0); // EOF indicates connection was closed
+
+  // The blocking read already synchronized - connection is definitely closed
+  CHECK(handler.connection_closed.load() ==
+        true); // Connection should be closed
+
+  close(fd);
+  server->shutdown();
+  runThread.join();
+}