my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import os
#2	import unittest
#3	from unittest.mock import patch
#4
#5	from embedchain import App
#6	from embedchain.config import AppConfig, ElasticsearchDBConfig
#7	from embedchain.embedder.gpt4all import GPT4AllEmbedder
#8	from embedchain.vectordb.elasticsearch import ElasticsearchDB
#9
#10
#11	class TestEsDB(unittest.TestCase):
#12	@patch("embedchain.vectordb.elasticsearch.Elasticsearch")
#13	def test_setUp(self, mock_client):
#14	self.db = ElasticsearchDB(config=ElasticsearchDBConfig(es_url="https://localhost:9200"))
#15	self.vector_dim = 384
#16	app_config = AppConfig(collect_metrics=False)
#17	self.app = App(config=app_config, db=self.db)
#18
#19	# Assert that the Elasticsearch client is stored in the ElasticsearchDB class.
#20	self.assertEqual(self.db.client, mock_client.return_value)
#21
#22	@patch("embedchain.vectordb.elasticsearch.Elasticsearch")
#23	def test_query(self, mock_client):
#24	self.db = ElasticsearchDB(config=ElasticsearchDBConfig(es_url="https://localhost:9200"))
#25	app_config = AppConfig(collect_metrics=False)
#26	self.app = App(config=app_config, db=self.db, embedding_model=GPT4AllEmbedder())
#27
#28	# Assert that the Elasticsearch client is stored in the ElasticsearchDB class.
#29	self.assertEqual(self.db.client, mock_client.return_value)
#30
#31	# Create some dummy data
#32	documents = ["This is a document.", "This is another document."]
#33	metadatas = [{"url": "url_1", "doc_id": "doc_id_1"}, {"url": "url_2", "doc_id": "doc_id_2"}]
#34	ids = ["doc_1", "doc_2"]
#35
#36	# Add the data to the database.
#37	self.db.add(documents, metadatas, ids)
#38
#39	search_response = {
#40	"hits": {
#41	"hits": [
#42	{
#43	"_source": {"text": "This is a document.", "metadata": {"url": "url_1", "doc_id": "doc_id_1"}},
#44	"_score": 0.9,
#45	},
#46	{
#47	"_source": {
#48	"text": "This is another document.",
#49	"metadata": {"url": "url_2", "doc_id": "doc_id_2"},
#50	},
#51	"_score": 0.8,
#52	},
#53	]
#54	}
#55	}
#56
#57	# Configure the mock client to return the mocked response.
#58	mock_client.return_value.search.return_value = search_response
#59
#60	# Query the database for the documents that are most similar to the query "This is a document".
#61	query = "This is a document"
#62	results_without_citations = self.db.query(query, n_results=2, where={})
#63	expected_results_without_citations = ["This is a document.", "This is another document."]
#64	self.assertEqual(results_without_citations, expected_results_without_citations)
#65
#66	results_with_citations = self.db.query(query, n_results=2, where={}, citations=True)
#67	expected_results_with_citations = [
#68	("This is a document.", {"url": "url_1", "doc_id": "doc_id_1", "score": 0.9}),
#69	("This is another document.", {"url": "url_2", "doc_id": "doc_id_2", "score": 0.8}),
#70	]
#71	self.assertEqual(results_with_citations, expected_results_with_citations)
#72
#73	def test_init_without_url(self):
#74	# Make sure it's not loaded from env
#75	try:
#76	del os.environ["ELASTICSEARCH_URL"]
#77	except KeyError:
#78	pass
#79	# Test if an exception is raised when an invalid es_config is provided
#80	with self.assertRaises(AttributeError):
#81	ElasticsearchDB()
#82
#83	def test_init_with_invalid_es_config(self):
#84	# Test if an exception is raised when an invalid es_config is provided
#85	with self.assertRaises(TypeError):
#86	ElasticsearchDB(es_config={"ES_URL": "some_url", "valid es_config": False})
#87

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public