etc/unittest/model.py


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

import unittest
import g4f
from g4f import ChatCompletion
from .mocks import ModelProviderMock

DEFAULT_MESSAGES = [{'role': 'user', 'content': 'Hello'}]
        
test_model = g4f.models.Model(
    name          = "test/test_model",
    base_provider = "",
    best_provider = ModelProviderMock
)
g4f.models.ModelUtils.convert["test_model"] = test_model
        
class TestPassModel(unittest.TestCase):

    def test_model_instance(self):
        response = ChatCompletion.create(test_model, DEFAULT_MESSAGES)
        self.assertEqual(test_model.name, response)
        
    def test_model_name(self):
        response = ChatCompletion.create("test_model", DEFAULT_MESSAGES)
        self.assertEqual(test_model.name, response)
        
    def test_model_pass(self):
        response = ChatCompletion.create("test/test_model", DEFAULT_MESSAGES, ModelProviderMock)
        self.assertEqual(test_model.name, response)