Skip to content

Commit 64f254b

Browse files
added unit test
1 parent 1f6192e commit 64f254b

File tree

1 file changed

+34
-0
lines changed

1 file changed

+34
-0
lines changed

tests/test_mindsdb/test_knowledgebase.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from mindsdb_sql_parser.ast.mindsdb.knowledge_base import (
33
CreateKnowledgeBase,
44
DropKnowledgeBase,
5+
EvaluateKnowledgeBase
56
)
67
from mindsdb_sql_parser.ast import (
78
Select,
@@ -370,3 +371,36 @@ def test_insert_into_knowledge_base(self):
370371
),
371372
)
372373
assert ast == expected_ast
374+
375+
def test_evaluate_knowledge_base(self):
376+
sql = """
377+
EVALUATE my_knowledge_base
378+
USING
379+
TEST_TABLE = my_database.some_table_1,
380+
SAVE_TO = my_database.some_table_2,
381+
LLM = {
382+
"provider": "openai",
383+
"model": "gpt-3.5-turbo",
384+
"api_key": "my_api_key"
385+
},
386+
generate_data = {
387+
"from_sql": "SELECT content FROM my_database.some_table",
388+
"count": 100
389+
}
390+
"""
391+
ast = parse_sql(sql)
392+
expected_ast = EvaluateKnowledgeBase(
393+
name=Identifier("my_knowledge_base"),
394+
test_table=Identifier(parts=["my_database", "some_table_1"]),
395+
save_to=Identifier(parts=["my_database", "some_table_2"]),
396+
llm={
397+
"provider": "openai",
398+
"model": "gpt-3.5-turbo",
399+
"api_key": "my_api_key"
400+
},
401+
generate_data={
402+
"from_sql": "SELECT content FROM my_database.some_table",
403+
"count": 100
404+
}
405+
)
406+
assert ast == expected_ast

0 commit comments

Comments
 (0)