// Copyright (c) 2026 Advanced Micro Devices, Inc.
// SPDX-License-Identifier: MIT

/// @file daemon_test.cpp
/// @brief Runs HIP kernel tests through the rocjitsu daemon via RPC.
///
/// Each test starts a daemon process, runs the HIP test binary with
/// LD_PRELOAD as a subprocess, verifies the result, and tears down the
/// daemon. Paths are injected via CMake compile definitions.

#include <gtest/gtest.h>

#include <array>
#include <cerrno>
#include <chrono>
#include <csignal>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <string>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/un.h>
#include <sys/wait.h>
#include <system_error>
#include <thread>
#include <unistd.h>
#include <vector>

namespace {

struct ProcessResult {
  std::string output;
  int exit_code = -1;
};

bool daemon_ready(const std::string &path) {
  int fd = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0);
  if (fd < 0)
    return false;
  sockaddr_un addr{};
  addr.sun_family = AF_UNIX;
  path.copy(addr.sun_path, sizeof(addr.sun_path) - 1);
  bool ok = connect(fd, reinterpret_cast<sockaddr *>(&addr), sizeof(addr)) == 0;
  close(fd);
  return ok;
}

struct TestPaths {
  std::string daemon_bin = RJ_DAEMON_BIN;
  std::string daemon_config = RJ_DAEMON_CONFIG;
  std::string daemon_config_2gpu = RJ_DAEMON_CONFIG_2GPU;
  std::string sanitizer_preload = RJ_DAEMON_SANITIZER_PRELOAD;
  std::string preload_lib = RJ_PRELOAD_LIB;
  std::string hip_vector_add_bin = RJ_HIP_VECTOR_ADD_BIN;
  std::string hip_memcpy_bin = RJ_HIP_MEMCPY_BIN;
  std::string hip_rccl_bin = RJ_HIP_RCCL_BIN;
  std::string daemon_logging_config = RJ_DAEMON_LOGGING_CONFIG;
  std::string interposer_dup_bin = RJ_INTERPOSER_DUP_BIN;
};

std::filesystem::path resolve_relative_to_exe(const std::filesystem::path &exe_dir,
                                              const char *path) {
  std::filesystem::path candidate(path);
  if (!candidate.is_absolute())
    candidate = exe_dir / candidate;
  return candidate.lexically_normal();
}

std::filesystem::path current_exe_dir() {
  std::error_code ec;
  std::filesystem::path exe = std::filesystem::read_symlink("/proc/self/exe", ec);
  if (ec)
    return {};
  return exe.parent_path();
}

bool installed_paths_exist(const TestPaths &paths) {
  // Only the UNCONDITIONALLY-built artifacts are required here. hip_rccl_bin is
  // deliberately excluded: the RCCL test binary (and the RcclDaemonTest cases
  // that use it) are built/registered only when RCCL_LIB is found at configure
  // time, so requiring it would make this check fail on non-RCCL installs and
  // wrongly fall back to build-tree paths that a pure install does not have.
  return std::filesystem::exists(paths.daemon_bin) &&
         std::filesystem::exists(paths.daemon_config) &&
         std::filesystem::exists(paths.daemon_config_2gpu) &&
         std::filesystem::exists(paths.preload_lib) &&
         std::filesystem::exists(paths.hip_vector_add_bin) &&
         std::filesystem::exists(paths.hip_memcpy_bin) &&
         std::filesystem::exists(paths.daemon_logging_config) &&
         std::filesystem::exists(paths.interposer_dup_bin);
}

TestPaths installed_paths(const std::filesystem::path &exe_dir) {
  return {
      resolve_relative_to_exe(exe_dir, RJ_INSTALLED_DAEMON_BIN).string(),
      resolve_relative_to_exe(exe_dir, RJ_INSTALLED_DAEMON_CONFIG).string(),
      resolve_relative_to_exe(exe_dir, RJ_INSTALLED_DAEMON_CONFIG_2GPU).string(),
      RJ_DAEMON_SANITIZER_PRELOAD,
      resolve_relative_to_exe(exe_dir, RJ_INSTALLED_PRELOAD_LIB).string(),
      resolve_relative_to_exe(exe_dir, RJ_INSTALLED_HIP_VECTOR_ADD_BIN).string(),
      resolve_relative_to_exe(exe_dir, RJ_INSTALLED_HIP_MEMCPY_BIN).string(),
      resolve_relative_to_exe(exe_dir, RJ_INSTALLED_HIP_RCCL_BIN).string(),
      resolve_relative_to_exe(exe_dir, RJ_INSTALLED_DAEMON_LOGGING_CONFIG).string(),
      resolve_relative_to_exe(exe_dir, RJ_INSTALLED_INTERPOSER_DUP_BIN).string(),
  };
}

TestPaths &test_paths() {
  static TestPaths paths = [] {
    std::filesystem::path exe_dir = current_exe_dir();
    if (!exe_dir.empty()) {
      TestPaths installed = installed_paths(exe_dir);
      if (installed_paths_exist(installed))
        return installed;
    }
    return TestPaths{};
  }();
  return paths;
}

const char *daemon_bin() { return test_paths().daemon_bin.c_str(); }

const char *daemon_config() { return test_paths().daemon_config.c_str(); }

const char *daemon_config_2gpu() { return test_paths().daemon_config_2gpu.c_str(); }

const char *sanitizer_preload() { return test_paths().sanitizer_preload.c_str(); }

const char *preload_lib() { return test_paths().preload_lib.c_str(); }

const char *hip_vector_add_bin() { return test_paths().hip_vector_add_bin.c_str(); }

const char *hip_memcpy_bin() { return test_paths().hip_memcpy_bin.c_str(); }

const char *hip_rccl_bin() { return test_paths().hip_rccl_bin.c_str(); }

const char *daemon_logging_config() { return test_paths().daemon_logging_config.c_str(); }
const char *interposer_dup_bin() { return test_paths().interposer_dup_bin.c_str(); }

TEST(RocjitsuCliDaemon, LaunchesApplicationAfterDaemonIsReady) {
  const char *xdg = std::getenv("XDG_RUNTIME_DIR");
  std::string tmp_dir = std::string(xdg ? xdg : "/tmp") + "/rocjitsu-launch-XXXXXX";
  ASSERT_NE(mkdtemp(tmp_dir.data()), nullptr) << "mkdtemp failed: " << strerror(errno);
  struct TempCleanup {
    std::string path;
    ~TempCleanup() {
      std::error_code error;
      std::filesystem::remove_all(path, error);
    }
  } cleanup{tmp_dir};
  const std::string runtime_dir = tmp_dir + "/rocjitsu";
  const std::string socket_path = runtime_dir + "/daemon.sock";

  const pid_t launcher = fork();
  ASSERT_GE(launcher, 0) << "fork failed: " << strerror(errno);
  if (launcher == 0) {
    setenv("XDG_RUNTIME_DIR", tmp_dir.c_str(), 1);
    setenv("ROCJITSU_RUNTIME_DIR", runtime_dir.c_str(), 1);
    execl(daemon_bin(), daemon_bin(), "--daemon", "--config", daemon_config(), "--",
          hip_vector_add_bin(), "--gtest_filter=HipVectorAddTest.CorrectResult", nullptr);
    _exit(127);
  }

  int status = 0;
  ASSERT_EQ(waitpid(launcher, &status, 0), launcher);
  EXPECT_TRUE(WIFEXITED(status));
  EXPECT_EQ(WEXITSTATUS(status), 0);

  const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5);
  while (std::filesystem::exists(socket_path) && std::chrono::steady_clock::now() < deadline)
    std::this_thread::sleep_for(std::chrono::milliseconds(20));
  EXPECT_FALSE(std::filesystem::exists(socket_path));
}

class DaemonTest : public ::testing::Test {
protected:
  // Config the daemon is launched with. The base fixture uses the plain KMD
  // config; subclasses can synthesize a config (e.g. to enable plugins) into
  // tmp_dir_ and return its path.
  virtual std::string daemon_config_for_test() { return daemon_config(); }

  void SetUp() override {
    // Each test gets its own unique subdirectory under XDG_RUNTIME_DIR
    // for the daemon socket, preventing conflicts when tests run in parallel.
    const char *xdg = std::getenv("XDG_RUNTIME_DIR");
    std::string base = xdg ? xdg : "/tmp";
    std::string tmpl = base + "/rocjitsu-test-XXXXXX";
    ASSERT_NE(mkdtemp(tmpl.data()), nullptr) << "mkdtemp failed: " << strerror(errno);
    tmp_dir_ = tmpl;
    runtime_dir_ = tmp_dir_ + "/rocjitsu";
    sock_path_ = runtime_dir_ + "/daemon.sock";

    ASSERT_NO_FATAL_FAILURE(config_path_ = daemon_config_for_test());
    ASSERT_FALSE(config_path_.empty()) << "daemon config unavailable";

    daemon_pid_ = fork();
    ASSERT_GE(daemon_pid_, 0) << "fork failed: " << strerror(errno);

    if (daemon_pid_ == 0) {
      setenv("XDG_RUNTIME_DIR", tmp_dir_.c_str(), 1);
      setenv("ROCJITSU_RUNTIME_DIR", runtime_dir_.c_str(), 1);
      execl(daemon_bin(), daemon_bin(), "--daemon", "--config", config_path_.c_str(), nullptr);
      _exit(127);
    }

    auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30);
    while (std::chrono::steady_clock::now() < deadline) {
      if (daemon_ready(sock_path_))
        return;
      int status = 0;
      if (waitpid(daemon_pid_, &status, WNOHANG) > 0) {
        daemon_pid_ = -1;
        FAIL() << "daemon exited before creating socket";
      }
      std::this_thread::sleep_for(std::chrono::milliseconds(50));
    }
    FAIL() << "daemon socket not created after 30s";
  }

  void TearDown() override {
    if (daemon_pid_ > 0) {
      EXPECT_EQ(kill(daemon_pid_, SIGTERM), 0);
      int status = 0;
      EXPECT_EQ(waitpid(daemon_pid_, &status, 0), daemon_pid_);
      EXPECT_TRUE(WIFEXITED(status));
      EXPECT_EQ(WEXITSTATUS(status), 0);
      EXPECT_FALSE(std::filesystem::exists(sock_path_));
      daemon_pid_ = -1;
    }
    std::filesystem::remove_all(tmp_dir_);
  }

  ProcessResult run_hip_test(const char *binary, const char *gtest_filter) {
    // CTest sets ROCJITSU_RUNTIME_DIR for isolation and the RPC layer prefers it
    // over XDG_RUNTIME_DIR. Override both here so the daemon and all client
    // subprocesses agree on the same socket and config-path directory.
    std::string cmd = "ROCJITSU_RUNTIME_DIR=";
    cmd += runtime_dir_;
    cmd += " XDG_RUNTIME_DIR=";
    cmd += tmp_dir_;
    cmd += " LD_PRELOAD=";
    cmd += sanitizer_preload();
    cmd += preload_lib();
    cmd += " HSA_ENABLE_SDMA=1 ";
    cmd += binary;
    if (gtest_filter && gtest_filter[0]) {
      cmd += " --gtest_filter=";
      cmd += gtest_filter;
    }
    cmd += " 2>&1";

    ProcessResult result;
    std::array<char, 4096> buf;
    FILE *pipe = popen(cmd.c_str(), "r");
    if (!pipe) {
      result.exit_code = -1;
      return result;
    }
    while (fgets(buf.data(), buf.size(), pipe) != nullptr)
      result.output += buf.data();
    int status = pclose(pipe);
    result.exit_code = WIFEXITED(status) ? WEXITSTATUS(status) : -1;
    return result;
  }

  ProcessResult run_rccl_rank(int rank, int world_size, const std::string &shared_dir,
                              const char *gtest_filter) {
    std::string cmd = "ROCJITSU_RUNTIME_DIR=";
    cmd += runtime_dir_;
    cmd += " XDG_RUNTIME_DIR=";
    cmd += tmp_dir_;
    cmd += " ";
    cmd += daemon_bin();
    cmd += " --attach --config ";
    cmd += daemon_config();
    cmd += " -- ";
    cmd += hip_rccl_bin();
    cmd += " --rank=";
    cmd += std::to_string(rank);
    cmd += " --world-size=";
    cmd += std::to_string(world_size);
    cmd += " --shared-dir=";
    cmd += shared_dir;
    if (gtest_filter && gtest_filter[0]) {
      cmd += " --gtest_filter=";
      cmd += gtest_filter;
    }
    cmd += " 2>&1";

    ProcessResult result;
    std::array<char, 4096> buf;
    FILE *pipe = popen(cmd.c_str(), "r");
    if (!pipe) {
      result.exit_code = -1;
      return result;
    }
    while (fgets(buf.data(), buf.size(), pipe) != nullptr)
      result.output += buf.data();
    int status = pclose(pipe);
    result.exit_code = WIFEXITED(status) ? WEXITSTATUS(status) : -1;
    return result;
  }

  void run_collective(const char *filter, int world_size = 2) {
    std::string shared_tmpl = tmp_dir_ + "/coll-XXXXXX";
    ASSERT_NE(mkdtemp(shared_tmpl.data()), nullptr) << strerror(errno);
    std::string shared_dir = shared_tmpl;

    std::vector<std::thread> threads(world_size);
    std::vector<ProcessResult> results(world_size);

    for (int r = 0; r < world_size; ++r)
      threads[r] =
          std::thread([&, r] { results[r] = run_rccl_rank(r, world_size, shared_dir, filter); });
    for (auto &t : threads)
      t.join();

    for (int r = 0; r < world_size; ++r)
      EXPECT_EQ(results[r].exit_code, 0) << "Rank " << r << " failed:\n" << results[r].output;

    std::filesystem::remove_all(shared_dir);
  }

  pid_t daemon_pid_ = -1;
  std::string tmp_dir_;
  std::string runtime_dir_;
  std::string sock_path_;
  std::string config_path_;
};

// --- hip_vector_add_test ---

TEST_F(DaemonTest, HipVectorAdd) {
  auto r = run_hip_test(hip_vector_add_bin(), "HipVectorAddTest.CorrectResult");
  EXPECT_EQ(r.exit_code, 0) << r.output;
}

TEST_F(DaemonTest, AttachExecFailurePreservesDaemonSocket) {
  const pid_t client = fork();
  ASSERT_GE(client, 0) << "fork failed: " << strerror(errno);
  if (client == 0) {
    setenv("XDG_RUNTIME_DIR", tmp_dir_.c_str(), 1);
    setenv("ROCJITSU_RUNTIME_DIR", runtime_dir_.c_str(), 1);
    execl(daemon_bin(), daemon_bin(), "--attach", "--config", daemon_config(), "--",
          "/does/not/exist", nullptr);
    _exit(127);
  }

  int status = 0;
  ASSERT_EQ(waitpid(client, &status, 0), client);
  EXPECT_TRUE(WIFEXITED(status));
  EXPECT_EQ(WEXITSTATUS(status), 1);
  EXPECT_TRUE(std::filesystem::is_socket(sock_path_));
  EXPECT_TRUE(daemon_ready(sock_path_));
}

// --- hip_memcpy_test ---

TEST_F(DaemonTest, HipMemcpyRoundTripFloat) {
  auto r = run_hip_test(hip_memcpy_bin(), "HipMemcpyTest.RoundTripFloat");
  EXPECT_EQ(r.exit_code, 0) << r.output;
}

TEST_F(DaemonTest, HipMemcpyRoundTripInt) {
  auto r = run_hip_test(hip_memcpy_bin(), "HipMemcpyTest.RoundTripInt");
  EXPECT_EQ(r.exit_code, 0) << r.output;
}

TEST_F(DaemonTest, HipMemcpyDeviceToDevice) {
  auto r = run_hip_test(hip_memcpy_bin(), "HipMemcpyTest.DeviceToDevice");
  EXPECT_EQ(r.exit_code, 0) << r.output;
}

// --- Multi-client tests ---

TEST_F(DaemonTest, TwoIndependentClients) {
  std::thread t1, t2;
  ProcessResult r1, r2;

  t1 = std::thread(
      [&] { r1 = run_hip_test(hip_vector_add_bin(), "HipVectorAddTest.CorrectResult"); });
  t2 = std::thread([&] { r2 = run_hip_test(hip_memcpy_bin(), "HipMemcpyTest.RoundTripFloat"); });

  t1.join();
  t2.join();

  EXPECT_EQ(r1.exit_code, 0) << "Client 1 (vector_add):\n" << r1.output;
  EXPECT_EQ(r2.exit_code, 0) << "Client 2 (memcpy):\n" << r2.output;
}

// --- Daemon-mode plugin coverage ---
//
// Plugins declared in the config load and run inside the forked daemon, not the
// interposer. The daemon resolves librocjitsu_plugin_logging.so by explicit
// path from the CLI's own directory, so this exercises the same plugin
// discovery + load contract as the interposer path.
class DaemonPluginTest : public DaemonTest {
protected:
  // Rewrite the logging-plugin config template's placeholder sink dir to a
  // private, writable directory for this test, then run the daemon on it.
  std::string daemon_config_for_test() override {
    sink_dir_ = tmp_dir_ + "/plugin_sink";
    std::error_code ec;
    std::filesystem::create_directories(sink_dir_, ec);

    std::ifstream in(daemon_logging_config());
    if (!in.good()) {
      ADD_FAILURE() << "cannot open daemon logging config: " << daemon_logging_config();
      return {};
    }
    std::string cfg{std::istreambuf_iterator<char>(in), std::istreambuf_iterator<char>()};

    const std::string token = RJ_DAEMON_LOGGING_SINK_TOKEN;
    for (size_t pos = cfg.find(token); pos != std::string::npos; pos = cfg.find(token, pos)) {
      cfg.replace(pos, token.size(), sink_dir_);
      pos += sink_dir_.size();
    }

    std::string out_path = tmp_dir_ + "/daemon_logging_config.json";
    std::ofstream out(out_path);
    out << cfg;
    if (!out.good()) {
      ADD_FAILURE() << "cannot write daemon logging config: " << out_path;
      return {};
    }
    return out_path;
  }

  std::string sink_dir_;
};

TEST_F(DaemonPluginTest, LoggingPluginDispatchLogged) {
  auto r = run_hip_test(hip_vector_add_bin(), "HipVectorAddTest.CorrectResult");
  ASSERT_EQ(r.exit_code, 0) << r.output;

  // The logging plugin runs in the daemon and writes <sink_dir>/logging.log as
  // it executes the client's kernel dispatch. Its presence and contents prove
  // the plugin loaded and ran in daemon mode.
  std::string log = sink_dir_ + "/logging.log";
  ASSERT_TRUE(std::filesystem::exists(log))
      << "daemon did not produce plugin log at " << log << "\n"
      << r.output;
  std::ifstream f(log);
  std::string contents{std::istreambuf_iterator<char>(f), std::istreambuf_iterator<char>()};
  EXPECT_NE(contents.find("dispatch"), std::string::npos) << "log:\n" << contents;
  EXPECT_NE(contents.find("vgprs="), std::string::npos) << "log:\n" << contents;
}

// --- Remote-backend dup/backend bookkeeping ---
//
// Runs the interposer dup/dup2/dup3 regression binary against a live daemon, so
// open("/dev/kfd") is serviced by the RemoteDriver (Remote backend) rather than
// the in-process SimulatedKfd. This gives the fd/backend state machine coverage
// on the remote path — the primary-fd re-mint (reissue_synthetic_kfd_fd) and the
// invalidation-vs-open serialization — which the CLI-launched (local) variant of
// the same tests cannot reach.

TEST_F(DaemonTest, InterposerDupReopenAfterPrimaryOverwriteRemote) {
  auto r = run_hip_test(interposer_dup_bin(),
                        "InterposerDupTest.ReopenAfterPrimaryOverwriteKeepsBackend");
  EXPECT_EQ(r.exit_code, 0) << r.output;
}

TEST_F(DaemonTest, InterposerDupSerializedReopenUnderContentionRemote) {
  auto r = run_hip_test(interposer_dup_bin(),
                        "InterposerDupTest.SerializedReopenUnderContentionStaysRoutable");
  EXPECT_EQ(r.exit_code, 0) << r.output;
}

TEST_F(DaemonTest, InterposerDupKeepsRoutingAfterPrimaryCloseRemote) {
  auto r =
      run_hip_test(interposer_dup_bin(), "InterposerDupTest.DupKeepsKfdRoutingAfterPrimaryClose");
  EXPECT_EQ(r.exit_code, 0) << r.output;
}

TEST_F(DaemonTest, InterposerDup2OverPrimaryInvalidatesRemote) {
  auto r =
      run_hip_test(interposer_dup_bin(), "InterposerDupTest.Dup2OverPrimaryInvalidatesKfdIdentity");
  EXPECT_EQ(r.exit_code, 0) << r.output;
}

// --- RCCL collective tests (2-GPU daemon) ---

class RcclDaemonTest : public ::testing::Test {
protected:
  void SetUp() override {
    const char *xdg = std::getenv("XDG_RUNTIME_DIR");
    std::string base = xdg ? xdg : "/tmp";
    std::string tmpl = base + "/rocjitsu-rccl-XXXXXX";
    ASSERT_NE(mkdtemp(tmpl.data()), nullptr) << "mkdtemp failed: " << strerror(errno);
    tmp_dir_ = tmpl;
    runtime_dir_ = tmp_dir_ + "/rocjitsu";
    sock_path_ = runtime_dir_ + "/daemon.sock";

    daemon_pid_ = fork();
    ASSERT_GE(daemon_pid_, 0) << "fork failed: " << strerror(errno);

    if (daemon_pid_ == 0) {
      setenv("XDG_RUNTIME_DIR", tmp_dir_.c_str(), 1);
      setenv("ROCJITSU_RUNTIME_DIR", runtime_dir_.c_str(), 1);
      execl(daemon_bin(), daemon_bin(), "--daemon", "--config", daemon_config_2gpu(), nullptr);
      _exit(127);
    }

    auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30);
    while (std::chrono::steady_clock::now() < deadline) {
      if (daemon_ready(sock_path_))
        return;
      int status = 0;
      if (waitpid(daemon_pid_, &status, WNOHANG) > 0) {
        daemon_pid_ = -1;
        FAIL() << "daemon exited before creating socket";
      }
      std::this_thread::sleep_for(std::chrono::milliseconds(50));
    }
    FAIL() << "daemon socket not created after 30s";
  }

  void TearDown() override {
    if (daemon_pid_ > 0) {
      EXPECT_EQ(kill(daemon_pid_, SIGTERM), 0);
      int status = 0;
      EXPECT_EQ(waitpid(daemon_pid_, &status, 0), daemon_pid_);
      EXPECT_TRUE(WIFEXITED(status));
      EXPECT_EQ(WEXITSTATUS(status), 0);
      EXPECT_FALSE(std::filesystem::exists(sock_path_));
      daemon_pid_ = -1;
    }
    std::filesystem::remove_all(tmp_dir_);
    std::this_thread::sleep_for(std::chrono::milliseconds(500));
  }

  ProcessResult run_rccl_rank(int rank, int world_size, const std::string &shared_dir,
                              const char *gtest_filter) {
    std::string cmd = "timeout 150 env ROCJITSU_RUNTIME_DIR=";
    cmd += runtime_dir_;
    cmd += " XDG_RUNTIME_DIR=";
    cmd += tmp_dir_;
    cmd += " HIP_VISIBLE_DEVICES=";
    cmd += std::to_string(rank);
    cmd += " NCCL_P2P_DISABLE=1 NCCL_SHM_DISABLE=1 HSA_NO_SCRATCH_RECLAIM=1"
           " NCCL_SOCKET_NTHREADS=1 NCCL_NSOCKS_PERTHREAD=1"
           " NCCL_SOCKET_IFNAME=lo ";
    cmd += daemon_bin();
    cmd += " --attach --config ";
    cmd += daemon_config_2gpu();
    cmd += " -- ";
    cmd += hip_rccl_bin();
    cmd += " --rank=";
    cmd += std::to_string(rank);
    cmd += " --world-size=";
    cmd += std::to_string(world_size);
    cmd += " --shared-dir=";
    cmd += shared_dir;
    if (gtest_filter && gtest_filter[0]) {
      cmd += " --gtest_filter=";
      cmd += gtest_filter;
    }
    cmd += " 2>&1";

    ProcessResult result;
    std::array<char, 4096> buf;
    FILE *pipe = popen(cmd.c_str(), "r");
    if (!pipe) {
      result.exit_code = -1;
      return result;
    }
    while (fgets(buf.data(), buf.size(), pipe) != nullptr)
      result.output += buf.data();
    int status = pclose(pipe);
    result.exit_code = WIFEXITED(status) ? WEXITSTATUS(status) : -1;
    return result;
  }

  void run_collective(const char *filter, int world_size = 2) {
    std::string shared_tmpl = tmp_dir_ + "/coll-XXXXXX";
    ASSERT_NE(mkdtemp(shared_tmpl.data()), nullptr) << strerror(errno);
    std::string shared_dir = shared_tmpl;

    std::vector<std::thread> threads(world_size);
    std::vector<ProcessResult> results(world_size);

    for (int r = 0; r < world_size; ++r)
      threads[r] =
          std::thread([&, r] { results[r] = run_rccl_rank(r, world_size, shared_dir, filter); });
    for (auto &t : threads)
      t.join();

    for (int r = 0; r < world_size; ++r)
      EXPECT_EQ(results[r].exit_code, 0) << "Rank " << r << " failed:\n" << results[r].output;

    std::filesystem::remove_all(shared_dir);
  }

  pid_t daemon_pid_ = -1;
  std::string tmp_dir_;
  std::string runtime_dir_;
  std::string sock_path_;
};

TEST_F(RcclDaemonTest, AllReduce) { run_collective("RcclTest.AllReduce"); }
TEST_F(RcclDaemonTest, Broadcast) { run_collective("RcclTest.Broadcast"); }
TEST_F(RcclDaemonTest, AllGather) { run_collective("RcclTest.AllGather"); }
TEST_F(RcclDaemonTest, ReduceScatter) { run_collective("RcclTest.ReduceScatter"); }
TEST_F(RcclDaemonTest, SendRecv) { run_collective("RcclTest.SendRecv"); }

} // namespace