daggy/daggyr/tests/unit_server.cpp

#include <curl/curl.h>
#include <rapidjson/document.h>
#include <sys/stat.h>

#include <catch2/catch.hpp>
#include <daggy/Serialization.hpp>
#include <daggy/executors/task/ForkingTaskExecutor.hpp>
#include <daggy/executors/task/NoopTaskExecutor.hpp>
#include <daggy/loggers/dag_run/OStreamLogger.hpp>
#include <daggyr/Server.hpp>
#include <filesystem>
#include <iostream>
#include <thread>

namespace rj = rapidjson;

using namespace daggy;

TEST_CASE("rest_endpoint", "[server_basic]")
{
  std::stringstream ss;
  Pistache::Address listenSpec("localhost", Pistache::Port(0));

  const ssize_t maxCores = 10, maxMemoryMB = 1000;

  daggyr::Server server(listenSpec, maxCores, maxMemoryMB);
  server.init(10);
  server.start();

  const std::string host    = "localhost:";
  const std::string baseURL = host + std::to_string(server.getPort());

  SECTION("Ready Endpoint")
  {
    auto response = HTTP_REQUEST(baseURL + "/ready");
    REQUIRE(response.code == HTTPCode::Ok);
  }

  SECTION("Querying a non-existent task should yield a 404")
  {
    auto response = HTTP_REQUEST(baseURL + "/v1/task/100/sample_name");
    REQUIRE(response.code == HTTPCode::Not_Found);
  }

  SECTION("Task Missing Cores should Fail")
  {
    std::string taskSpec =
        R"({ "job": { "command": [ "/usr/bin/touch", "dagrun_{{FILE}}" ]}, "memoryMB": 100 })";

    auto response =
        HTTP_REQUEST(baseURL + "/v1/task/0/sample_task", taskSpec, "POST");
    REQUIRE(response.code == HTTPCode::Not_Acceptable);
  }

  SECTION("Task Missing MemoryMB should Fail")
  {
    std::string taskSpec =
        R"({ "job": { "command": [ "/usr/bin/touch", "dagrun_{{FILE}}" ]}, "cores": 100 })";

    auto response =
        HTTP_REQUEST(baseURL + "/v1/task/0/sample_task", taskSpec, "POST");
    REQUIRE(response.code == HTTPCode::Not_Acceptable);
  }

  SECTION("Task submission and get result")
  {
    std::string taskSpec =
        R"({ "job": { "command": [ "/usr/bin/echo", "hello", "world" ], "cores": "1", "memoryMB": "100" } })";

    // Submit
    {
      auto response =
          HTTP_REQUEST(baseURL + "/v1/task/0/sample_task", taskSpec, "POST");
      REQUIRE(response.code == HTTPCode::Ok);
    }

    while (true) {
      auto [code, doc] = JSON_HTTP_REQUEST(baseURL + "/v1/task/0/sample_task");
      REQUIRE(doc.IsObject());
      REQUIRE(doc.HasMember("state"));

      std::string state = doc["state"].GetString();
      if (state != "COMPLETED") {
        std::this_thread::sleep_for(250ms);
      }
      else {
        REQUIRE(doc.HasMember("attempt"));
        auto attempt = attemptRecordFromJSON(doc["attempt"]);

        REQUIRE(attempt.rc == 0);
        REQUIRE(attempt.outputLog == "hello world\n");
        break;
      }
    }
  }

  SECTION("Task capacity changes")
  {
    std::string taskSpec =
        R"({ "job": { "command": [ "/usr/bin/sleep", "5" ], "cores": "1", "memoryMB": "100" } })";

    auto getCapacity = [&]() -> daggy::executors::task::daggy_runner::Capacity {
      daggy::executors::task::daggy_runner::Capacity cap;
      auto [code, doc] = JSON_HTTP_REQUEST(baseURL + "/v1/capacity");
      REQUIRE(doc.IsObject());
      REQUIRE(doc.HasMember("current"));
      const auto &cur = doc["current"];
      REQUIRE(cur.IsObject());
      REQUIRE(cur.HasMember("cores"));
      REQUIRE(cur.HasMember("memoryMB"));

      cap.cores    = cur["cores"].GetInt64();
      cap.memoryMB = cur["memoryMB"].GetInt64();

      return cap;
    };

    auto preCap = getCapacity();

    // Submit
    {
      auto response =
          HTTP_REQUEST(baseURL + "/v1/task/0/sample_task", taskSpec, "POST");
      REQUIRE(response.code == HTTPCode::Ok);
    }

    auto postCap = getCapacity();

    REQUIRE(postCap.cores == preCap.cores - 1);
    REQUIRE(postCap.memoryMB == preCap.memoryMB - 100);

    // Ensure the current job is running
    {
      auto [code, doc] = JSON_HTTP_REQUEST(baseURL + "/v1/task/0/sample_task");
      REQUIRE(doc.IsObject());
      REQUIRE(doc.HasMember("state"));
      REQUIRE(doc["state"] != "COMPLETED");
    }

    // Stop it
    {
      auto [code, doc] =
          JSON_HTTP_REQUEST(baseURL + "/v1/task/0/sample_task", "", "DELETE");
      REQUIRE(code == HTTPCode::Ok);
    }

    // Grab it and ensure it was killed
    while (true) {
      auto response = HTTP_REQUEST(baseURL + "/v1/task/0/sample_task");

      REQUIRE(response.code == HTTPCode::Ok);
      rj::Document doc;
      daggy::checkRJParse(doc.Parse(response.body.c_str()));
      REQUIRE(doc.IsObject());
      REQUIRE(doc.HasMember("state"));

      std::string state = doc["state"].GetString();
      if (state != "COMPLETED") {
        std::this_thread::sleep_for(250ms);
      }
      else {
        REQUIRE(doc.HasMember("attempt"));
        auto attempt = attemptRecordFromJSON(doc["attempt"]);

        REQUIRE(attempt.rc != 0);
        break;
      }
    }
  }

  server.shutdown();
}