mirror of https://github.com/arc53/DocsGPT
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
33 lines
1.2 KiB
Python
33 lines
1.2 KiB
Python
import unittest
|
|
from unittest.mock import patch
|
|
from application.llm.openai import OpenAILLM
|
|
|
|
class TestOpenAILLM(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
self.api_key = "test_api_key"
|
|
self.llm = OpenAILLM(self.api_key)
|
|
|
|
def test_init(self):
|
|
self.assertEqual(self.llm.api_key, self.api_key)
|
|
|
|
@patch('application.llm.openai.openai.ChatCompletion.create')
|
|
def test_gen(self, mock_create):
|
|
model = "test_model"
|
|
engine = "test_engine"
|
|
messages = ["test_message"]
|
|
response = {"choices": [{"message": {"content": "test_response"}}]}
|
|
mock_create.return_value = response
|
|
result = self.llm.gen(model, engine, messages)
|
|
self.assertEqual(result, "test_response")
|
|
|
|
@patch('application.llm.openai.openai.ChatCompletion.create')
|
|
def test_gen_stream(self, mock_create):
|
|
model = "test_model"
|
|
engine = "test_engine"
|
|
messages = ["test_message"]
|
|
response = [{"choices": [{"delta": {"content": "test_response"}}]}]
|
|
mock_create.return_value = response
|
|
result = list(self.llm.gen_stream(model, engine, messages))
|
|
self.assertEqual(result, ["test_response"])
|