Update test-suites.

zh-plus · Jul 28, 2024 · d55fb5b · d55fb5b
1 parent f1f39c5
commit d55fb5b
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 10 deletions.
diff --git a/tests/test_chatbot.py b/tests/test_chatbot.py
@@ -40,7 +40,7 @@ def test_estimate_fee(self):
             {'role': 'user', 'content': 'Hello'},
         ]
         fee = bot.estimate_fee(messages)
-        self.assertTrue(isclose(fee, 6e-06))
+        self.assertTrue(isclose(fee, 2.1e-06))
 
     def test_gpt_update_fee(self):
         bot = self.gpt_bot
@@ -55,7 +55,7 @@ def test_gpt_update_fee(self):
         bot.api_fees += [0]
         response3 = OpenAIResponse(usage=OpenAIUsage(prompt_tokens=300, completion_tokens=600, total_tokens=900))
         bot.update_fee(response3)
-        self.assertListEqual(bot.api_fees, [0.00035, 0.0007, 0.00105])
+        self.assertListEqual(bot.api_fees, [0.000135, 0.00027, 0.000405])
 
     def test_claude_update_fee(self):
         bot = self.claude_bot

diff --git a/tests/test_openlrc.py b/tests/test_openlrc.py
@@ -13,16 +13,16 @@
 from openlrc.utils import extend_filename
 
 
-@patch('openlrc.transcribe.WhisperModel', MagicMock())
+@patch('openlrc.transcribe.BatchedInferencePipeline', MagicMock())
 @patch('openlrc.transcribe.Transcriber.transcribe',
        MagicMock(return_value=(
                [
                    Segment(
-                       0, 0, 0, 3, 'hello world1', [], 0, 0.8, 0, 0, words=[
+                       0, 0, 0, 3, 'hello world1', [], 0.8, 0, 0, words=[
                            Word(0, 1.5, 'hello', probability=0.8), Word(1.6, 3, ' world1', probability=0.8)
                        ]),
                    Segment(
-                       0, 0, 3, 6, 'hello world2', [], 0, 0.8, 0, 0, words=[
+                       0, 0, 3, 6, 'hello world2', [], 0.8, 0, 0, words=[
                            Word(3, 4.5, 'hello', probability=0.8), Word(4.6, 6, ' world2', probability=0.8)
                        ])
                ],

diff --git a/tests/test_transcribe.py b/tests/test_transcribe.py
@@ -11,12 +11,12 @@
 
 return_tuple = ([
                     Segment(
-                        0, 0, 0, 3, 'hello world', [], 0, 0.8, 0, 0, words=[
+                        0, 0, 0, 3, 'hello world', [], 0.8, 0, 0, words=[
                             Word(0, 1.5, 'hello', probability=0.8),
                             Word(1.6, 3, ' world', probability=0.8)
                         ]),
                     Segment(
-                        0, 0, 3, 6, 'hello world', [], 0, 0.8, 0, 0, words=[
+                        0, 0, 3, 6, 'hello world', [], 0.8, 0, 0, words=[
                             Word(3, 4.5, 'hello', probability=0.8),
                             Word(4.6, 6, ' world', probability=0.8)
                         ]),
@@ -28,7 +28,7 @@ def setUp(self) -> None:
         self.audio_path = Path('data/test_audio.wav')
 
     def test_transcribe_success(self):
-        with patch('openlrc.transcribe.WhisperModel') as MockModel:
+        with patch('openlrc.transcribe.BatchedInferencePipeline') as MockModel:
             MockModel.return_value.transcribe.return_value = return_tuple
 
             transcriber = Transcriber(model_name='tiny')
@@ -37,7 +37,7 @@ def test_transcribe_success(self):
             self.assertEqual(round(info.duration), 30)
 
     def test_audio_file_not_found(self):
-        with patch('openlrc.transcribe.WhisperModel') as MockModel:
+        with patch('openlrc.transcribe.BatchedInferencePipeline') as MockModel:
             MockModel.return_value.transcribe.return_value = return_tuple
 
             transcriber = Transcriber(model_name='tiny')

diff --git a/tests/test_translate.py b/tests/test_translate.py
@@ -10,7 +10,7 @@
 from openlrc.translate import LLMTranslator
 from openlrc.utils import get_similarity
 
-test_models = ['gpt-3.5-turbo', 'claude-3-haiku-20240307']
+test_models = ['gpt-4o-mini', 'claude-3-haiku-20240307']
 
 
 class TestLLMTranslator(unittest.TestCase):