arbiterAI/tests/arbiterAITests.cpp at main · caseymcc/arbiterAI · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
#include "arbiterAI/arbiterAI.h"
#include "arbiterAI/modelManager.h"
#include <gtest/gtest.h>
#include <gmock/gmock.h>

namespace arbiterAI
{

class ArbiterAITest : public ::testing::Test
{
protected:
    void SetUp() override
    {
        ModelManager::reset();
    }
};

//TEST_F(ArbiterAITest, GetDownloadStatus)
//{
//    arbiterAI ai;
//    ai.initialize({});
//    std::string error;
//
//    // Initialize ModelManager so the model can be found.
//    ModelManager &mm = ModelManager::instance();
//    ModelInfo model;
//    model.model = "non-existent-model";
//    model.provider = "llama";
//    mm.addModel(model);
//
//    EXPECT_EQ(ai.getDownloadStatus("non-existent-model", error), DownloadStatus::NotStarted);
//}

TEST_F(ArbiterAITest, SupportModelDownload)
{
    ArbiterAI ai;
    ModelManager &mm = ModelManager::instance();

    ModelInfo localModel;
    localModel.model = "local-model";
    localModel.provider = "llama"; // Llama provider supports downloads
    mm.addModel(localModel);

    ModelInfo remoteModel;
    remoteModel.model = "remote-model";
    remoteModel.provider = "openai"; // OpenAI provider does not support downloads
    mm.addModel(remoteModel);

    EXPECT_TRUE(ai.supportModelDownload("llama"));
    EXPECT_FALSE(ai.supportModelDownload("openai"));
    EXPECT_FALSE(ai.supportModelDownload("unknown-provider"));
}

} // namespace arbiterAI