wehub-resource-sync
65 lines
2.3 KiB
Python
65 lines
2.3 KiB
Python
""" Tests paths both for a standard account setup and custom account setup.
|
|
|
|
Note: by default, this test will run on MongoDB - to switch DB used - add one-line:
|
|
|
|
`LLMWareConfig().set_active_db("sqlite")
|
|
|
|
"""
|
|
|
|
import os
|
|
|
|
from llmware.configs import LLMWareConfig
|
|
from llmware.library import Library
|
|
from llmware.retrieval import Query
|
|
from llmware.setup import Setup
|
|
|
|
|
|
def test_standard_account_setup():
|
|
|
|
library_name = "test_standard_account_setup"
|
|
library = Library().create_new_library(library_name)
|
|
|
|
# add files
|
|
sample_files_path = Setup().load_sample_files()
|
|
library.add_files(os.path.join(sample_files_path,"SmallLibrary"))
|
|
|
|
# run test query
|
|
results = Query(library).text_query("base salary")
|
|
|
|
assert len(results) > 0, "Expected results to be greater than 0 for standard account setup."
|
|
|
|
|
|
# Test overriding llmware_data folder
|
|
def test_setup_custom_location():
|
|
|
|
# upstream application creating custom home path, e.g., aibloks
|
|
new_home_folder = "llmware_data_custom"
|
|
LLMWareConfig.set_home(os.path.join(LLMWareConfig.get_home(), new_home_folder))
|
|
|
|
new_llmware_path = "aibloks_data"
|
|
LLMWareConfig.set_llmware_path_name(new_llmware_path)
|
|
|
|
# Make folders if they don't already exist
|
|
if not os.path.exists(LLMWareConfig().get_home()):
|
|
os.mkdir(LLMWareConfig().get_home())
|
|
if not os.path.exists(LLMWareConfig().get_llmware_path()):
|
|
LLMWareConfig.setup_llmware_workspace()
|
|
|
|
assert new_home_folder in LLMWareConfig.get_home(), "Custom home folder not set correctly."
|
|
assert new_home_folder in LLMWareConfig.get_llmware_path() and new_llmware_path in LLMWareConfig.get_llmware_path(), "Custom LLMWare path not set correctly."
|
|
assert new_home_folder in LLMWareConfig.get_library_path() and new_llmware_path in LLMWareConfig.get_library_path(), "Library path not set correctly."
|
|
|
|
# test with library setup and adding files
|
|
library_name = "library_with_custom_path"
|
|
library = Library().create_new_library(library_name)
|
|
sample_files_path = Setup().load_sample_files()
|
|
library.add_files(os.path.join(sample_files_path,"Agreements"))
|
|
|
|
assert new_home_folder in library.library_main_path, "Custom path not found in library main path."
|
|
|
|
results = Query(library).text_query("salary")
|
|
|
|
assert len(results) > 0
|
|
|
|
|