From c27e559f10864f54647215b64cdb3f11b5eb6192 Mon Sep 17 00:00:00 2001
From: binary-husky <qingxu.fu@outlook.com>
Date: Tue, 6 Feb 2024 17:51:47 +0800
Subject: [PATCH] match sess-* key

---
 shared_utils/key_pattern_manager.py |  2 +-
 tests/test_key_pattern_manager.py   | 41 +++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_key_pattern_manager.py

diff --git a/shared_utils/key_pattern_manager.py b/shared_utils/key_pattern_manager.py
index eed2af9..a617b7f 100644
--- a/shared_utils/key_pattern_manager.py
+++ b/shared_utils/key_pattern_manager.py
@@ -14,7 +14,7 @@ def is_openai_api_key(key):
     if len(CUSTOM_API_KEY_PATTERN) != 0:
         API_MATCH_ORIGINAL = re.match(CUSTOM_API_KEY_PATTERN, key)
     else:
-        API_MATCH_ORIGINAL = re.match(r"sk-[a-zA-Z0-9]{48}$", key)
+        API_MATCH_ORIGINAL = re.match(r"sk-[a-zA-Z0-9]{48}$|sess-[a-zA-Z0-9]{40}$", key)
     return bool(API_MATCH_ORIGINAL)
 
 
diff --git a/tests/test_key_pattern_manager.py b/tests/test_key_pattern_manager.py
new file mode 100644
index 0000000..bf84441
--- /dev/null
+++ b/tests/test_key_pattern_manager.py
@@ -0,0 +1,41 @@
+import unittest
+
+def validate_path():
+    import os, sys
+
+    os.path.dirname(__file__)
+    root_dir_assume = os.path.abspath(os.path.dirname(__file__) + "/..")
+    os.chdir(root_dir_assume)
+    sys.path.append(root_dir_assume)
+
+
+validate_path()  # validate path so you can run from base directory
+
+from shared_utils.key_pattern_manager import is_openai_api_key
+
+class TestKeyPatternManager(unittest.TestCase):
+    def test_is_openai_api_key_with_valid_key(self):
+        key = "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
+        self.assertTrue(is_openai_api_key(key))
+
+        key = "sx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
+        self.assertFalse(is_openai_api_key(key))
+
+        key = "sess-wg61ZafYHpNz7FFwIH7HGZlbVqUVaeV5tatHCWpl"
+        self.assertTrue(is_openai_api_key(key))
+
+        key = "sess-wg61ZafYHpNz7FFwIH7HGZlbVqUVa5tatHCWpl"
+        self.assertFalse(is_openai_api_key(key))
+
+
+    def test_is_openai_api_key_with_invalid_key(self):
+        key = "invalid_key"
+        self.assertFalse(is_openai_api_key(key))
+
+    def test_is_openai_api_key_with_custom_pattern(self):
+        # Assuming you have set a custom pattern in your configuration
+        key = "custom-pattern-key"
+        self.assertFalse(is_openai_api_key(key))
+
+if __name__ == '__main__':
+    unittest.main()
\ No newline at end of file