From 05a18d60199bc40036ce007eabf90992d5be8ce1 Mon Sep 17 00:00:00 2001
From: Atomie CHEN <atomic_cwh@163.com>
Date: Tue, 17 Sep 2024 04:26:50 +0800
Subject: [PATCH] test(OpenAIClient): add tests for ensure_client_credentials &
 endpoint_manager

---
 tests/test_client.py | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/tests/test_client.py b/tests/test_client.py
index 1b21bef..e446b83 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -154,6 +154,8 @@ def test_ensure_client_credentials():
         client.chat(messages=[], api_key="should_not_be_used").api_key == "client_key"
     )
 
+
+def test_ensure_client_credentials2():
     client2 = OpenAIClient(ensure_client_credentials=True, api_key="client_key")
     assert (
         client2.chat(messages=[], api_key="should_not_be_used").api_key == "client_key"
@@ -162,3 +164,36 @@ def test_ensure_client_credentials():
     assert (
         client2.chat(messages=[], api_key="should_be_used").api_key == "should_be_used"
     )
+
+
+def test_ensure_client_credentials3():
+    client3 = OpenAIClient(
+        ensure_client_credentials=True,
+        endpoints=[
+            {
+                "api_key": "client_key1",
+            },
+            {
+                "api_key": "client_key2",
+            },
+            {
+                "api_key": "client_key3",
+            },
+        ],
+    )
+    assert (
+        client3.chat(messages=[], api_key="should_not_be_used").api_key == "client_key1"
+    )
+    assert (
+        client3.chat(messages=[], api_key="should_not_be_used").api_key == "client_key2"
+    )
+    assert (
+        client3.chat(messages=[], api_key="should_not_be_used").api_key == "client_key3"
+    )
+    assert (
+        client3.chat(messages=[], api_key="should_not_be_used").api_key == "client_key1"
+    )
+    client3.ensure_client_credentials = False
+    assert (
+        client3.chat(messages=[], api_key="should_be_used").api_key == "should_be_used"
+    )