yams/tests/plugins/onnx_provider_plugin_catch2_test.cpp

#include <nlohmann/json.hpp>
#include <filesystem>
#include <fstream>
#include <string>
#include "../common/env_compat.h"
#include "../common/test_helpers_catch2.h"
#include <catch2/catch_test_macros.hpp>
#include <yams/compat/dlfcn.h>
#include <yams/compat/unistd.h>

extern "C" {
#include <yams/plugins/abi.h>
#include <yams/plugins/model_provider_v1.h>
}

static yams_model_provider_v1* load_provider(void** handle) {
#ifndef ONNX_PLUGIN_PATH
#ifdef _WIN32
#define ONNX_PLUGIN_PATH "build/debug/plugins/onnx/yams_onnx_plugin.dll"
#elif defined(__APPLE__)
#define ONNX_PLUGIN_PATH "build/debug/plugins/onnx/libyams_onnx_plugin.dylib"
#else
#define ONNX_PLUGIN_PATH "build/debug/plugins/onnx/libyams_onnx_plugin.so"
#endif
#endif
    const char* so = ONNX_PLUGIN_PATH;
    *handle = dlopen(so, RTLD_LAZY | RTLD_LOCAL);
    if (!*handle)
        return nullptr;
    auto getiface =
        (int (*)(const char*, uint32_t, void**))dlsym(*handle, "yams_plugin_get_interface");
    if (!getiface)
        return nullptr;
    void* p = nullptr;
    int rc = getiface(YAMS_IFACE_MODEL_PROVIDER_V1, 2, &p);
    if (rc != 0 || !p)
        return nullptr;
    return reinterpret_cast<yams_model_provider_v1*>(p);
}

static int get_health_json(void* handle, std::string& out_json) {
    if (!handle)
        return -1;
    auto get_health = (int (*)(char**))dlsym(handle, "yams_plugin_get_health_json");
    if (!get_health)
        return -2;
    char* s = nullptr;
    int rc = get_health(&s);
    if (rc == 0 && s) {
        out_json.assign(s);
        std::free(s);
    }
    return rc;
}

TEST_CASE("OnnxProviderPluginTest.RuntimeInfoIncludesHints", "[plugin][onnxproviderplugintest]") {
    yams::test::ScopedEnvVar skipLoading{"YAMS_SKIP_MODEL_LOADING", std::string{"1"}};
    void* h = nullptr;
    auto* prov = load_provider(&h);
    if (!prov)
        SKIP("onnx plugin not available");
    char* js = nullptr;
    REQUIRE((0 == prov->get_runtime_info_json(prov->self, "e5-small", &js)));
    REQUIRE((js != nullptr));
    auto j = nlohmann::json::parse(js, nullptr, false);
    prov->free_string(prov->self, js);
    REQUIRE_FALSE(j.is_discarded());
    CHECK((j.value("model", "") == "e5-small"));
    CHECK(j.contains("graph_optimization"));
    CHECK(j.contains("execution_provider"));
}

TEST_CASE("OnnxProviderPluginTest.MissingRuntimeReportedViaHealth",
          "[plugin][onnxproviderplugintest]") {
    yams::test::ScopedEnvVar runtimeLibrary{"YAMS_ONNX_RUNTIME_LIB",
                                            std::string{"/nonexistent/libonnxruntime.dylib"}};
    yams::test::ScopedEnvVar skipLoading{"YAMS_SKIP_MODEL_LOADING", std::string{"1"}};

    void* h = nullptr;
    auto* prov = load_provider(&h);
    if (!prov)
        SKIP("onnx plugin not available");

    std::string health;
    REQUIRE((0 == get_health_json(h, health)));
    auto j = nlohmann::json::parse(health, nullptr, false);
    REQUIRE_FALSE(j.is_discarded());
    CHECK((j.value("status", "unknown") == "unavailable"));
    CHECK((j.value("reason", "") == "runtime_load_failed" ||
           j.value("reason", "") == "runtime_not_found"));
    CHECK_FALSE(j.value("runtime_status", "").empty());
    CHECK_FALSE(j.value("runtime_error", "").empty());
}

TEST_CASE("OnnxProviderPluginTest.ConfigPreloadParsingAndReady",
          "[plugin][onnxproviderplugintest]") {
    // Create temp XDG config with preload sections
    namespace fs = std::filesystem;
    fs::path tmp = fs::temp_directory_path() / "yams_test_cfg";
    fs::create_directories(tmp / "yams");
    std::ofstream(tmp / "yams" / "config.toml")
        << "[plugins.onnx]\npreload=\"a,b\"\n\n[plugins.onnx.models.sentiment]\n"
        << "task=\"sentiment\"\n";
    yams::test::ScopedEnvVar configHome{"XDG_CONFIG_HOME", tmp.string()};
    yams::test::ScopedEnvVar skipLoading{"YAMS_SKIP_MODEL_LOADING", std::string{"1"}};

    void* h = nullptr;
    auto* prov = load_provider(&h);
    if (!prov)
        SKIP("onnx plugin not available");
    std::string health;
    REQUIRE((0 == get_health_json(h, health)));
    auto j = nlohmann::json::parse(health, nullptr, false);
    REQUIRE_FALSE(j.is_discarded());
    CHECK((j.value("status", "unknown") != "unavailable"));
    // Prefer ok status
    CHECK((j.value("status", "ok") == "ok"));
}

TEST_CASE("OnnxProviderPluginTest.GenerateEmbeddingBatchErrors",
          "[plugin][onnxproviderplugintest]") {
    yams::test::ScopedEnvVar skipLoading{"YAMS_SKIP_MODEL_LOADING", std::string{"1"}};
    void* h = nullptr;
    auto* prov = load_provider(&h);
    if (!prov)
        SKIP("onnx plugin not available");

    // Invalid args: null inputs
    float* out = nullptr;
    size_t b = 0, d = 0;
    CHECK((0 != prov->generate_embedding_batch(prov->self, "bad-model", nullptr, nullptr, 0, &out,
                                               &b, &d)));

    // Bad model id with empty batch arrays
    const uint8_t* inputs_dummy = nullptr;
    size_t lens_dummy = 0; // not dereferenced for batch_size=0
    CHECK((0 != prov->generate_embedding_batch(prov->self, "bad-model", &inputs_dummy, &lens_dummy,
                                               0, &out, &b, &d)));
}