daggy/tests/unit_server.cpp

#include <iostream>
#include <filesystem>
#include <fstream>

#include <catch2/catch.hpp>
#include <pistache/client.h>
#include <rapidjson/document.h>

#include <daggy/Server.hpp>
#include <daggy/Serialization.hpp>
#include <daggy/executors/task/ForkingTaskExecutor.hpp>
#include <daggy/loggers/dag_run/OStreamLogger.hpp>

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"].GetArray();
            REQUIRE(taskStates.Size() == 3);

            complete = true;
            for (size_t i = 0; i < taskStates.Size(); ++i) {
                std::string state = taskStates[i].GetString();
                if (state != "COMPLETED") {
                    complete = false;
                    break;
                }
            }
            if (complete) break;
            std::this_thread::sleep_for(std::chrono::milliseconds(500));
        }
        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();
}