Katon Minhas
7724a7ddc4
Draft: update unit tests * update unit tests * Update unit tests * Update unit tests Approved-by: Alex Galarce
65 lines
2.9 KiB
Python
65 lines
2.9 KiB
Python
import unittest
|
|
from unittest.mock import patch, mock_open
|
|
import hashlib
|
|
import csv
|
|
from collections import OrderedDict
|
|
import src.utils.llm_utils as llm_utils
|
|
from src import config
|
|
|
|
class TestLLMUtils(unittest.TestCase):
|
|
def setUp(self):
|
|
self.filename = "test_file.txt"
|
|
self.prompt = "This is a test prompt."
|
|
self.response = "This is a test response."
|
|
self.model_type = "test_model"
|
|
self.input_cost_per_1k = 0.01
|
|
self.output_cost_per_1k = 0.02
|
|
self.mock_stats = {self.filename: {"requests": 0, "tokens_sent": 0, "tokens_received": 0, "total_time": 0}}
|
|
config.MODEL_STATS = self.mock_stats
|
|
config.GLOBAL_STATS = {}
|
|
|
|
def test_update_stats(self):
|
|
llm_utils.update_stats(self.filename, 10, 20, 5, self.model_type)
|
|
self.assertEqual(config.MODEL_STATS[self.filename]["requests"], 1)
|
|
self.assertEqual(config.MODEL_STATS[self.filename]["tokens_sent"], 10)
|
|
self.assertEqual(config.MODEL_STATS[self.filename]["tokens_received"], 20)
|
|
self.assertEqual(config.MODEL_STATS[self.filename]["total_time"], 5)
|
|
self.assertEqual(config.GLOBAL_STATS["total_requests"], 1)
|
|
self.assertEqual(config.GLOBAL_STATS["total_tokens_sent"], 10)
|
|
self.assertEqual(config.GLOBAL_STATS["total_tokens_received"], 20)
|
|
self.assertEqual(config.GLOBAL_STATS["total_time"], 5)
|
|
|
|
def test_count_tokens(self):
|
|
self.assertEqual(llm_utils.count_tokens("123456"), 1)
|
|
self.assertEqual(llm_utils.count_tokens("1234567"), 2)
|
|
self.assertEqual(llm_utils.count_tokens(""), 0)
|
|
|
|
def test_create_cost_log_entry(self):
|
|
entry = llm_utils.create_cost_log_entry(
|
|
self.filename, "caller_func", self.prompt, self.response, self.input_cost_per_1k, self.output_cost_per_1k
|
|
)
|
|
self.assertEqual(entry["Filename"], self.filename)
|
|
self.assertEqual(entry["Caller Function"], "caller_func")
|
|
self.assertEqual(entry["Input Prompt Length"], len(self.prompt))
|
|
self.assertEqual(entry["Output Response Length"], len(self.response))
|
|
self.assertIn("Total Cost", entry)
|
|
|
|
def test_get_cache_key(self):
|
|
model_id = "test_model"
|
|
expected_hash = hashlib.sha256(f"{model_id}:{self.prompt}".encode()).hexdigest()
|
|
self.assertEqual(llm_utils.get_cache_key(self.prompt, model_id), expected_hash)
|
|
|
|
@patch("builtins.open", new_callable=mock_open)
|
|
@patch("csv.DictWriter.writerow")
|
|
def test_invoke_claude_local_mode(self, mock_csv_writer, mock_open_file):
|
|
config.RUN_MODE = "local"
|
|
model_id = config.MODEL_ID_CLAUDE2
|
|
|
|
with patch("src.utils.llm_utils.local_claude_2", return_value=self.response):
|
|
response = llm_utils.invoke_claude(self.prompt, model_id, self.filename)
|
|
self.assertEqual(response, self.response)
|
|
mock_csv_writer.assert_called()
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|