update

KenyonY · Apr 30, 2023 · 67af3e7 · 67af3e7
1 parent 8291899
commit 67af3e7
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/openai_forward/_base.py b/openai_forward/_base.py
@@ -16,8 +16,12 @@ class OpenaiBase:
     IP_BLACKLIST = os.environ.get("IP_BLACKLIST", "").strip()
     if IP_BLACKLIST:
         IP_BLACKLIST = [i.strip() for i in IP_BLACKLIST.split(' ')]
+    else:
+        IP_BLACKLIST = []
     if IP_WHITELIST:
         IP_WHITELIST = [i.strip() for i in IP_WHITELIST.split(' ')]
+    else:
+        IP_WHITELIST = []
     if _ROUTE_PREFIX:
         if _ROUTE_PREFIX.endswith('/'):
             _ROUTE_PREFIX = _ROUTE_PREFIX[:-1]

diff --git a/tests/test_api.py b/tests/test_api.py
@@ -1,4 +1,5 @@
 from openai_forward.openai import OpenaiBase
+from fastapi import HTTPException
 import pytest
 
 
@@ -14,11 +15,13 @@ def test_env(self, openai: OpenaiBase):
     def test_validate_ip(self, openai: OpenaiBase):
         ip1 = "1.1.1.1"
         ip2 = "2.2.2.2"
-        assert openai.validate_request_host("*")
+        assert openai.validate_request_host("*") is None
         openai.IP_WHITELIST.append(ip1)
-        assert openai.validate_request_host(ip1)
-        assert openai.validate_request_host(ip2) is False
+        assert openai.validate_request_host(ip1) is None
+        with pytest.raises(HTTPException):
+            openai.validate_request_host(ip2)
         openai.IP_WHITELIST = []
         openai.IP_BLACKLIST.append(ip1)
-        assert openai.validate_request_host(ip1) is False
-        assert openai.validate_request_host(ip2)
+        assert openai.validate_request_host(ip2) is None
+        with pytest.raises(HTTPException):
+            openai.validate_request_host(ip1)