49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
import argparse
|
|
import unittest
|
|
from unittest.mock import MagicMock
|
|
from chatmastermind.ai_factory import create_ai
|
|
from chatmastermind.configuration import Config
|
|
from chatmastermind.ai import AIError
|
|
from chatmastermind.ais.openai import OpenAI
|
|
|
|
|
|
class TestCreateAI(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.args = MagicMock(spec=argparse.Namespace)
|
|
self.args.AI = 'myopenai'
|
|
self.args.model = None
|
|
self.args.max_tokens = None
|
|
self.args.temperature = None
|
|
|
|
def test_create_ai_from_args(self) -> None:
|
|
# Create an AI with the default configuration
|
|
config = Config()
|
|
self.args.AI = 'myopenai'
|
|
ai = create_ai(self.args, config)
|
|
self.assertIsInstance(ai, OpenAI)
|
|
|
|
def test_create_ai_from_default(self) -> None:
|
|
self.args.AI = None
|
|
# Create an AI with the default configuration
|
|
config = Config()
|
|
ai = create_ai(self.args, config)
|
|
self.assertIsInstance(ai, OpenAI)
|
|
|
|
def test_create_empty_ai_error(self) -> None:
|
|
self.args.AI = None
|
|
# Create Config with empty AIs
|
|
config = Config()
|
|
config.ais = {}
|
|
# Call create_ai function and assert that it raises AIError
|
|
with self.assertRaises(AIError):
|
|
create_ai(self.args, config)
|
|
|
|
def test_create_unsupported_ai_error(self) -> None:
|
|
# Mock argparse.Namespace with ai='invalid_ai'
|
|
self.args.AI = 'invalid_ai'
|
|
# Create default Config
|
|
config = Config()
|
|
# Call create_ai function and assert that it raises AIError
|
|
with self.assertRaises(AIError):
|
|
create_ai(self.args, config)
|