From d733930a2b1876340039d90f19ece81fab0d078d Mon Sep 17 00:00:00 2001 From: Heiner Lohaus Date: Fri, 23 Feb 2024 02:51:10 +0100 Subject: Fix unittests, use Union typing --- etc/unittest/client.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) (limited to 'etc') diff --git a/etc/unittest/client.py b/etc/unittest/client.py index 2bc00c2e..ec8aa4b7 100644 --- a/etc/unittest/client.py +++ b/etc/unittest/client.py @@ -35,13 +35,15 @@ class TestPassModel(unittest.TestCase): response = client.chat.completions.create(messages, "Hello", stream=True) for chunk in response: self.assertIsInstance(chunk, ChatCompletionChunk) - self.assertIsInstance(chunk.choices[0].delta.content, str) + if chunk.choices[0].delta.content is not None: + self.assertIsInstance(chunk.choices[0].delta.content, str) messages = [{'role': 'user', 'content': chunk} for chunk in ["You ", "You ", "Other", "?"]] response = client.chat.completions.create(messages, "Hello", stream=True, max_tokens=2) response = list(response) - self.assertEqual(len(response), 2) + self.assertEqual(len(response), 3) for chunk in response: - self.assertEqual(chunk.choices[0].delta.content, "You ") + if chunk.choices[0].delta.content is not None: + self.assertEqual(chunk.choices[0].delta.content, "You ") def test_stop(self): client = Client(provider=YieldProviderMock) -- cgit v1.2.3