daggy/tests/unit_server.cpp

#include <pistache/client.h>
#include <rapidjson/document.h>

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

namespace rj = rapidjson;

Pistache::Http::Response REQUEST(std::string url, std::string payload = "")
{
  Pistache::Http::Experimental::Client client;
  client.init();
  Pistache::Http::Response response;
  auto reqSpec = (payload.empty() ? client.get(url) : client.post(url));
  reqSpec.timeout(std::chrono::seconds(2));
  if (!payload.empty()) {
    reqSpec.body(payload);
  }
  auto request = reqSpec.send();
  bool ok = false, error = false;
  std::string msg;
  request.then(
      [&](Pistache::Http::Response rsp) {
        ok       = true;
        response = rsp;
      },
      [&](std::exception_ptr ptr) {
        error = true;
        try {
          std::rethrow_exception(ptr);
        }
        catch (std::exception &e) {
          msg = e.what();
        }
      });

  Pistache::Async::Barrier<Pistache::Http::Response> barrier(request);
  barrier.wait_for(std::chrono::seconds(2));
  client.shutdown();
  if (error) {
    throw std::runtime_error(msg);
  }
  return response;
}

TEST_CASE("rest_endpoint", "[server_basic]")
{
  std::stringstream ss;
  daggy::executors::task::ForkingTaskExecutor executor(10);
  daggy::loggers::dag_run::OStreamLogger logger(ss);
  Pistache::Address listenSpec("localhost", Pistache::Port(0));

  const size_t nDAGRunners = 10, nWebThreads = 10;

  daggy::Server server(listenSpec, logger, executor, nDAGRunners);
  server.init(nWebThreads);
  server.start();

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

  SECTION("Ready Endpoint")
  {
    auto response = REQUEST(baseURL + "/ready");
    REQUIRE(response.code() == Pistache::Http::Code::Ok);
  }

  SECTION("Querying a non-existent dagrunid should fail ")
  {
    auto response = REQUEST(baseURL + "/v1/dagrun/100");
    REQUIRE(response.code() != Pistache::Http::Code::Ok);
  }

  SECTION("Simple DAGRun Submission")
  {
    std::string dagRun = R"({
            "name": "unit_server",
            "parameters": { "FILE": [ "A", "B" ] },
            "tasks": {
               "touch": { "job": { "command": [ "/usr/bin/touch", "dagrun_{{FILE}}" ]} },
               "cat": {  "job": { "command": [ "/usr/bin/cat", "dagrun_A", "dagrun_B" ]},
                 "parents": [ "touch" ]
               }
               }
        })";

    // Submit, and get the runID
    daggy::DAGRunID runID = 0;
    {
      auto response = REQUEST(baseURL + "/v1/dagrun/", dagRun);
      REQUIRE(response.code() == Pistache::Http::Code::Ok);

      rj::Document doc;
      daggy::checkRJParse(doc.Parse(response.body().c_str()));
      REQUIRE(doc.IsObject());
      REQUIRE(doc.HasMember("runID"));

      runID = doc["runID"].GetUint64();
    }

    // Ensure our runID shows up in the list of running DAGs
    {
      auto response = REQUEST(baseURL + "/v1/dagrun/");
      REQUIRE(response.code() == Pistache::Http::Code::Ok);

      rj::Document doc;
      daggy::checkRJParse(doc.Parse(response.body().c_str()));
      REQUIRE(doc.IsArray());
      REQUIRE(doc.Size() >= 1);

      // Ensure that our DAG is in the list and matches our given DAGRunID
      bool found       = false;
      const auto &runs = doc.GetArray();
      for (size_t i = 0; i < runs.Size(); ++i) {
        const auto &run = runs[i];
        REQUIRE(run.IsObject());
        REQUIRE(run.HasMember("name"));
        REQUIRE(run.HasMember("runID"));

        std::string runName = run["name"].GetString();
        if (runName == "unit_server") {
          REQUIRE(run["runID"].GetUint64() == runID);
          found = true;
          break;
        }
      }
      REQUIRE(found);
    }

    // Wait until our DAG is complete
    bool complete = true;
    for (auto i = 0; i < 10; ++i) {
      auto response = REQUEST(baseURL + "/v1/dagrun/" + std::to_string(runID));
      REQUIRE(response.code() == Pistache::Http::Code::Ok);
      rj::Document doc;
      daggy::checkRJParse(doc.Parse(response.body().c_str()));
      REQUIRE(doc.IsObject());

      REQUIRE(doc.HasMember("taskStates"));
      const auto &taskStates = doc["taskStates"].GetObject();

      size_t nStates = 0;
      for (auto it = taskStates.MemberBegin(); it != taskStates.MemberEnd();
           ++it) {
        nStates++;
      }
      REQUIRE(nStates == 3);

      complete = true;
      for (auto it = taskStates.MemberBegin(); it != taskStates.MemberEnd();
           ++it) {
        std::string state = it->value.GetString();
        if (state != "COMPLETED") {
          complete = false;
          break;
        }
      }
      if (complete)
        break;
      std::this_thread::sleep_for(std::chrono::seconds(1));
    }
    REQUIRE(complete);

    std::this_thread::sleep_for(std::chrono::seconds(2));
    for (const auto &pth : std::vector<fs::path>{"dagrun_A", "dagrun_B"}) {
      REQUIRE(fs::exists(pth));
      fs::remove(pth);
    }
  }

  server.shutdown();
}