test: update test script

Author: jameszyao

test/testcase/test_async/test_async_assistant.py

@@ -3,23 +3,22 @@
 from taskingai.inference import *
 from test.config import embedding_model_id, chat_completion_model_id
 from test.common.logger import logger
+import re
 
 
 @pytest.mark.test_async
 class TestChatCompletion:
-
     @pytest.mark.run(order=4)
     @pytest.mark.asyncio
     async def test_a_chat_completion(self):
-
         # normal chat completion.
 
         normal_res = await a_chat_completion(
             model_id=chat_completion_model_id,
             messages=[
                 SystemMessage("You are a professional assistant."),
                 UserMessage("Hi"),
-            ]
+            ],
         )
         pytest.assume(normal_res.finish_reason == "stop")
         pytest.assume(normal_res.message.content)
@@ -35,10 +34,7 @@ async def test_a_chat_completion(self):
                 UserMessage("Hi"),
                 AssistantMessage("Hello! How can I assist you today?"),
                 UserMessage("Can you tell me a joke?"),
-                AssistantMessage(
-                    "Sure, here is a joke for you: Why don't scientists trust atoms? Because they make up everything!"),
-                UserMessage("That's funny. Can you tell me another one?"),
-            ]
+            ],
         )
 
         pytest.assume(multi_round_res.finish_reason == "stop")
@@ -55,13 +51,8 @@ async def test_a_chat_completion(self):
                 UserMessage("Hi"),
                 AssistantMessage("Hello! How can I assist you today?"),
                 UserMessage("Can you tell me a joke?"),
-                AssistantMessage(
-                    "Sure, here is a joke for you: Why don't scientists trust atoms? Because they make up everything!"),
-                UserMessage("That's funny. Can you tell me another one?"),
             ],
-            configs={
-                "max_tokens": 10
-            }
+            configs={"max_tokens": 10},
         )
         pytest.assume(max_tokens_res.finish_reason == "length")
         pytest.assume(max_tokens_res.message.content)
@@ -70,33 +61,34 @@ async def test_a_chat_completion(self):
 
         # chat completion with stream.
 
-        stream_res = await a_chat_completion(model_id=chat_completion_model_id,
-                                             messages=[
-                                                 SystemMessage("You are a professional assistant."),
-                                                 UserMessage("count from 1 to 50 and separate numbers by comma."),
-                                             ],
-                                             stream=True
-                                             )
-        except_list = [i + 1 for i in range(50)]
-        real_list = []
+        stream_res = await a_chat_completion(
+            model_id=chat_completion_model_id,
+            messages=[
+                SystemMessage("You are a professional assistant."),
+                UserMessage("count from 1 to 10 and separate numbers by comma."),
+            ],
+            stream=True,
+        )
+        except_list = [i + 1 for i in range(10)]
+        real_str = ""
         async for item in stream_res:
             if isinstance(item, ChatCompletionChunk):
                 logger.info(f"Message: {item.delta}")
-                if item.delta.isdigit():
-                    real_list.append(int(item.delta))
+                real_str += item.delta
+
             elif isinstance(item, ChatCompletion):
                 logger.info(f"Message: {item.finish_reason}")
                 pytest.assume(item.finish_reason == "stop")
+
+        real_list = [int(num) for num in re.findall(r"\b\d+\b", real_str)]
         pytest.assume(set(except_list) == set(real_list))
 
 
 @pytest.mark.test_async
 class TestTextEmbedding:
-
     @pytest.mark.run(order=0)
     @pytest.mark.asyncio
     async def test_a_text_embedding(self):
-
         # Text embedding with str.
 
         input_str = "Machine learning is a subfield of artificial intelligence (AI) that involves the development of algorithms that allow computers to learn from and make decisions or predictions based on data."

test/testcase/test_async/test_async_inference.py

@@ -22,8 +22,6 @@ class TestCollection(Base):
         "status",
     ]
     collection_keys = set(collection_list)
-    collection_configs = ["metric", "chunk_size", "chunk_overlap"]
-    collection_configs_keys = set(collection_configs)
 
     @pytest.mark.run(order=9)
     @pytest.mark.asyncio
@@ -43,7 +41,7 @@ async def test_a_create_collection(self):
             pytest.assume(res_dict["description"] == description)
             pytest.assume(res_dict["embedding_model_id"] == embedding_model_id)
             pytest.assume(res_dict["capacity"] == 1000)
-            pytest.assume(res_dict["status"] == "creating")
+            pytest.assume((res_dict["status"] == "ready") or (res_dict["status"] == "creating"))
 
     @pytest.mark.run(order=10)
     @pytest.mark.asyncio
@@ -76,8 +74,7 @@ async def test_a_get_collection(self, a_collection_id):
         res = await a_get_collection(collection_id=self.collection_id)
         res_dict = res.to_dict()
         pytest.assume(res_dict.keys() == self.collection_keys)
-        pytest.assume(res_dict["configs"].keys() == self.collection_configs_keys)
-        pytest.assume(res_dict["status"] == "ready" or "creating")
+        pytest.assume(res_dict["status"] == "ready" or res_dict["status"] == "creating")
 
     @pytest.mark.run(order=12)
     @pytest.mark.asyncio
@@ -89,7 +86,6 @@ async def test_a_update_collection(self):
         res = await a_update_collection(collection_id=self.collection_id, name=name, description=description)
         res_dict = res.to_dict()
         pytest.assume(res_dict.keys() == self.collection_keys)
-        pytest.assume(res_dict["configs"].keys() == self.collection_configs_keys)
         pytest.assume(res_dict["name"] == name)
         pytest.assume(res_dict["description"] == description)
         pytest.assume(res_dict["status"] == "ready")
@@ -99,17 +95,22 @@ async def test_a_update_collection(self):
     async def test_a_delete_collection(self):
         # List collections.
         old_res = await a_list_collections(order="desc", limit=100, after=None, before=None)
+        old_nums = len(old_res)
 
         for index, collection in enumerate(old_res):
             collection_id = collection.collection_id
-            # Delete a collection.
+            # Delete a collection
             await a_delete_collection(collection_id=collection_id)
 
             new_collections = await a_list_collections(order="desc", limit=100, after=None, before=None)
-            # List collections.
+
+            # List collections
             collection_ids = [c.collection_id for c in new_collections]
             pytest.assume(collection_id not in collection_ids)
 
+            new_nums = len(new_collections)
+            pytest.assume(new_nums == old_nums - 1 - index)
+
 
 @pytest.mark.test_async
 class TestRecord(Base):
@@ -125,8 +126,6 @@ class TestRecord(Base):
         "status",
     ]
     record_keys = set(record_list)
-    record_content = ["text"]
-    record_content_keys = set(record_content)
 
     @pytest.mark.run(order=13)
     @pytest.mark.asyncio
@@ -142,9 +141,8 @@ async def test_a_create_record(self):
             )
             res_dict = res.to_dict()
             pytest.assume(res_dict.keys() == self.record_keys)
-            pytest.assume(res_dict["content"].keys() == self.record_content_keys)
-            pytest.assume(res_dict["content"]["text"] == text)
-            pytest.assume(res_dict["status"] == "creating")
+            pytest.assume(res_dict["content"] == text)
+            pytest.assume((res_dict["status"] == "creating") or (res_dict["status"] == "ready"))
 
     @pytest.mark.run(order=14)
     @pytest.mark.asyncio
@@ -177,12 +175,14 @@ async def test_a_list_records(self, a_record_id):
     async def test_a_get_record(self):
         # Get a record.
 
-        res = await a_get_record(collection_id=self.collection_id, record_id=self.record_id)
-        logger.info(f"a_get_record:{res}")
-        res_dict = res.to_dict()
-        pytest.assume(res_dict.keys() == self.record_keys)
-        pytest.assume(res_dict["content"].keys() == self.record_content_keys)
-        pytest.assume(res_dict["status"] == "ready" or "creating")
+        records = await a_list_records(collection_id=self.collection_id)
+        for record in records:
+            record_id = record.record_id
+            res = await a_get_record(collection_id=self.collection_id, record_id=record_id)
+            logger.info(f"a_get_record:{res}")
+            res_dict = res.to_dict()
+            pytest.assume(res_dict.keys() == self.record_keys)
+            pytest.assume(res_dict["status"] == "ready" or res_dict["status"] == "creating")
 
     @pytest.mark.run(order=16)
     @pytest.mark.asyncio
@@ -194,7 +194,6 @@ async def test_a_update_record(self):
         logger.info(f"a_update_record:{res}")
         res_dict = res.to_dict()
         pytest.assume(res_dict.keys() == self.record_keys)
-        pytest.assume(res_dict["content"].keys() == self.record_content_keys)
         pytest.assume(res_dict["metadata"] == metadata)
 
     @pytest.mark.run(order=34)
@@ -226,7 +225,7 @@ async def test_a_delete_record(self):
 
 @pytest.mark.test_async
 class TestChunk(Base):
-    chunk_list = ["chunk_id", "collection_id", "record_id", "object", "text", "score"]
+    chunk_list = ["chunk_id", "collection_id", "record_id", "object", "content", "score"]
     chunk_keys = set(chunk_list)
 
     @pytest.mark.run(order=17)
@@ -241,5 +240,5 @@ async def test_a_query_chunks(self):
         for chunk in res:
             chunk_dict = chunk.to_dict()
             pytest.assume(chunk_dict.keys() == self.chunk_keys)
-            pytest.assume(query_text in chunk_dict["text"])
+            pytest.assume(query_text in chunk_dict["content"])
             pytest.assume(chunk_dict["score"] >= 0)

test/testcase/test_async/test_async_retrieval.py

@@ -26,7 +26,7 @@ class TestAction(Base):
     action_keys = set(action_list)
     action_schema = ["openapi", "info", "servers", "paths", "components", "security"]
     action_schema_keys = set(action_schema)
-    schema = {
+    openapi_schema = {
         "openapi": "3.1.0",
         "info": {
             "title": "Get weather data",
@@ -72,7 +72,7 @@ class TestAction(Base):
     async def test_a_bulk_create_actions(self):
         # Create an action.
 
-        res = await a_bulk_create_actions(schema=self.schema)
+        res = await a_bulk_create_actions(openapi_schema=self.openapi_schema)
         for action in res:
             action_dict = action.to_dict()
             logger.info(action_dict)
@@ -84,15 +84,15 @@ async def test_a_bulk_create_actions(self):
                     if action_dict["openapi_schema"][key]["/location"] == "get":
                         pytest.assume(
                             action_dict["openapi_schema"][key]["/location"]["get"]
-                            == self.schema["paths"]["/location"]["get"]
+                            == self.openapi_schema["paths"]["/location"]["get"]
                         )
                     elif action_dict["openapi_schema"][key]["/location"] == "post":
                         pytest.assume(
                             action_dict["openapi_schema"][key]["/location"]["post"]
-                            == self.schema["paths"]["/location"]["post"]
+                            == self.openapi_schema["paths"]["/location"]["post"]
                         )
                 else:
-                    pytest.assume(action_dict["openapi_schema"][key] == self.schema[key])
+                    pytest.assume(action_dict["openapi_schema"][key] == self.openapi_schema[key])
 
     @pytest.mark.run(order=5)
     @pytest.mark.asyncio
@@ -101,10 +101,10 @@ async def test_a_run_action(self, a_action_id):
 
         if not Base.action_id:
             Base.action_id = await a_action_id
-        parameters = {"location": "beijing"}
+        parameters = {"location": "tokyo"}
         res = await a_run_action(action_id=self.action_id, parameters=parameters)
         logger.info(f"async run action{res}")
-        pytest.assume(res["status"] == 400)
+        pytest.assume(res["status"] != 200)
         pytest.assume(res["error"])
 
     @pytest.mark.run(order=6)
@@ -176,7 +176,7 @@ async def test_a_update_action(self):
             "security": [],
         }
 
-        res = await a_update_action(action_id=self.action_id, schema=update_schema)
+        res = await a_update_action(action_id=self.action_id, openapi_schema=update_schema)
         res_dict = res.to_dict()
         logger.info(res_dict)
         pytest.assume(res_dict.keys() == self.action_keys)