test: add test for api key

privatellm/main.py

@@ -1,11 +1,11 @@
 """Program to play with langchain"""
 
-
 import glob
 import logging
 import os
 import tempfile
 from enum import Enum
+from pathlib import Path
 from timeit import default_timer as timer
 from typing import Any
 from urllib.parse import urljoin
@@ -121,11 +121,10 @@ def assert_api_key() -> str:
     api_key = os.environ.get("OPENAI_API_KEY", None)
     if api_key is not None:
         return api_key
-    if os.path.exists("apikey.txt"):
-        with open("apikey.txt", encoding="utf-8") as f:
+    path = Path("apikey.txt")
+    if path.exists():
+        with path.open(encoding="utf-8") as f:
             api_key = f.read().strip()
-    if not api_key:
-        api_key = os.getenv("OPENAI_API_KEY")
     if api_key:
         os.environ["OPENAI_API_KEY"] = api_key
         return api_key

tests/test_api_key.py

@@ -0,0 +1,30 @@
+import unittest
+from io import StringIO
+from pathlib import Path
+from unittest.mock import call, patch
+
+from privatellm.main import assert_api_key
+
+
+class ApiKeyTest(unittest.TestCase):
+    def test_empty(self):
+        with patch.dict("os.environ", {}), patch("privatellm.main.Path", spec_set=Path) as MockPath:
+            MockPath.return_value = MockPath
+            MockPath.exists.return_value = False
+            try:
+                assert_api_key()
+                self.fail("An assertion should have be thrown")
+            except AssertionError as exp:
+                assert exp.args == ("Missing OpenAI API key.",)
+            assert MockPath.call_args_list == [call("apikey.txt")]
+
+    def test_env(self):
+        with patch.dict("os.environ", {"OPENAI_API_KEY": "fooenv"}):
+            assert assert_api_key() == "fooenv"
+
+    def test_file(self):
+        with patch.dict("os.environ", {}), patch("privatellm.main.Path", spec_set=Path) as MockPath:
+            MockPath.return_value = MockPath
+            MockPath.exists.return_value = True
+            MockPath.open.return_value = StringIO(" foofile  ")  # expect stripped value
+            assert assert_api_key() == "foofile"