Update tests

Author: aravind-segu

databricks/sdk/mixins/open_ai_client.py

@@ -15,7 +15,6 @@ def __init__(self, get_headers_func):
                 self.get_headers_func = get_headers_func
 
             def auth_flow(self, request: httpx.Request) -> httpx.Request:
-                print("Calling Authenticate")
                 auth_headers = self.get_headers_func()
                 request.headers["Authorization"] = auth_headers["Authorization"]
                 yield request

tests/test_open_ai_mixin.py

@@ -11,10 +11,10 @@ def test_open_ai_client(monkeypatch):
     monkeypatch.setenv('DATABRICKS_HOST', 'test_host')
     monkeypatch.setenv('DATABRICKS_TOKEN', 'test_token')
     w = WorkspaceClient(config=Config())
-    w.serving_endpoints.get_open_ai_client()
+    client = w.serving_endpoints.get_open_ai_client()
 
-    # assert client.base_url == "https://test_host/serving-endpoints/"
-    # assert client.api_key == "test_token"
+    assert client.base_url == "https://test_host/serving-endpoints/"
+    assert client.api_key == "no-token"
 
 
 @pytest.mark.skipif(sys.version_info < (3, 8), reason="Requires Python > 3.7")
@@ -24,7 +24,7 @@ def test_langchain_open_ai_client(monkeypatch):
     monkeypatch.setenv('DATABRICKS_HOST', 'test_host')
     monkeypatch.setenv('DATABRICKS_TOKEN', 'test_token')
     w = WorkspaceClient(config=Config())
-    w.serving_endpoints.get_langchain_chat_open_ai_client("databricks-meta-llama-3-1-70b-instruct")
+    client = w.serving_endpoints.get_langchain_chat_open_ai_client("databricks-meta-llama-3-1-70b-instruct")
 
-    # assert client.openai_api_base == "https://test_host/serving-endpoints"
-    # assert client.model_name == "databricks-meta-llama-3-1-70b-instruct"
+    assert client.openai_api_base == "https://test_host/serving-endpoints"
+    assert client.model_name == "databricks-meta-llama-3-1-70b-instruct"