my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import os
#2	import shutil
#3
#4	import pytest
#5
#6	from embedchain import App
#7	from embedchain.config import AppConfig
#8	from embedchain.config.vector_db.lancedb import LanceDBConfig
#9	from embedchain.vectordb.lancedb import LanceDB
#10
#11	os.environ["OPENAI_API_KEY"] = "test-api-key"
#12
#13
#14	@pytest.fixture
#15	def lancedb():
#16	return LanceDB(config=LanceDBConfig(dir="test-db", collection_name="test-coll"))
#17
#18
#19	@pytest.fixture
#20	def app_with_settings():
#21	lancedb_config = LanceDBConfig(allow_reset=True, dir="test-db-reset")
#22	lancedb = LanceDB(config=lancedb_config)
#23	app_config = AppConfig(collect_metrics=False)
#24	return App(config=app_config, db=lancedb)
#25
#26
#27	@pytest.fixture(scope="session", autouse=True)
#28	def cleanup_db():
#29	yield
#30	try:
#31	shutil.rmtree("test-db.lance")
#32	shutil.rmtree("test-db-reset.lance")
#33	except OSError as e:
#34	print("Error: %s - %s." % (e.filename, e.strerror))
#35
#36
#37	def test_lancedb_duplicates_throw_warning(caplog):
#38	db = LanceDB(config=LanceDBConfig(allow_reset=True, dir="test-db"))
#39	app = App(config=AppConfig(collect_metrics=False), db=db)
#40	app.db.add(ids=["0"], documents=["doc1"], metadatas=["test"])
#41	app.db.add(ids=["0"], documents=["doc1"], metadatas=["test"])
#42	assert "Insert of existing doc ID: 0" not in caplog.text
#43	assert "Add of existing doc ID: 0" not in caplog.text
#44	app.db.reset()
#45
#46
#47	def test_lancedb_duplicates_collections_no_warning(caplog):
#48	db = LanceDB(config=LanceDBConfig(allow_reset=True, dir="test-db"))
#49	app = App(config=AppConfig(collect_metrics=False), db=db)
#50	app.set_collection_name("test_collection_1")
#51	app.db.add(ids=["0"], documents=["doc1"], metadatas=["test"])
#52	app.set_collection_name("test_collection_2")
#53	app.db.add(ids=["0"], documents=["doc1"], metadatas=["test"])
#54	assert "Insert of existing doc ID: 0" not in caplog.text
#55	assert "Add of existing doc ID: 0" not in caplog.text
#56	app.db.reset()
#57	app.set_collection_name("test_collection_1")
#58	app.db.reset()
#59
#60
#61	def test_lancedb_collection_init_with_default_collection():
#62	db = LanceDB(config=LanceDBConfig(allow_reset=True, dir="test-db"))
#63	app = App(config=AppConfig(collect_metrics=False), db=db)
#64	assert app.db.collection.name == "embedchain_store"
#65
#66
#67	def test_lancedb_collection_init_with_custom_collection():
#68	db = LanceDB(config=LanceDBConfig(allow_reset=True, dir="test-db"))
#69	app = App(config=AppConfig(collect_metrics=False), db=db)
#70	app.set_collection_name(name="test_collection")
#71	assert app.db.collection.name == "test_collection"
#72
#73
#74	def test_lancedb_collection_set_collection_name():
#75	db = LanceDB(config=LanceDBConfig(allow_reset=True, dir="test-db"))
#76	app = App(config=AppConfig(collect_metrics=False), db=db)
#77	app.set_collection_name("test_collection")
#78	assert app.db.collection.name == "test_collection"
#79
#80
#81	def test_lancedb_collection_changes_encapsulated():
#82	db = LanceDB(config=LanceDBConfig(allow_reset=True, dir="test-db"))
#83	app = App(config=AppConfig(collect_metrics=False), db=db)
#84	app.set_collection_name("test_collection_1")
#85	assert app.db.count() == 0
#86	app.db.add(ids=["0"], documents=["doc1"], metadatas=["test"])
#87	assert app.db.count() == 1
#88
#89	app.set_collection_name("test_collection_2")
#90	assert app.db.count() == 0
#91
#92	app.db.add(ids=["0"], documents=["doc1"], metadatas=["test"])
#93	app.set_collection_name("test_collection_1")
#94	assert app.db.count() == 1
#95	app.db.reset()
#96	app.set_collection_name("test_collection_2")
#97	app.db.reset()
#98
#99
#100	def test_lancedb_collection_collections_are_persistent():
#101	db = LanceDB(config=LanceDBConfig(allow_reset=True, dir="test-db"))
#102	app = App(config=AppConfig(collect_metrics=False), db=db)
#103	app.set_collection_name("test_collection_1")
#104	app.db.add(ids=["0"], documents=["doc1"], metadatas=["test"])
#105	del app
#106
#107	db = LanceDB(config=LanceDBConfig(allow_reset=True, dir="test-db"))
#108	app = App(config=AppConfig(collect_metrics=False), db=db)
#109	app.set_collection_name("test_collection_1")
#110	assert app.db.count() == 1
#111
#112	app.db.reset()
#113
#114
#115	def test_lancedb_collection_parallel_collections():
#116	db1 = LanceDB(config=LanceDBConfig(allow_reset=True, dir="test-db", collection_name="test_collection_1"))
#117	app1 = App(
#118	config=AppConfig(collect_metrics=False),
#119	db=db1,
#120	)
#121	db2 = LanceDB(config=LanceDBConfig(allow_reset=True, dir="test-db", collection_name="test_collection_2"))
#122	app2 = App(
#123	config=AppConfig(collect_metrics=False),
#124	db=db2,
#125	)
#126
#127	# cleanup if any previous tests failed or were interrupted
#128	app1.db.reset()
#129	app2.db.reset()
#130
#131	app1.db.add(ids=["0"], documents=["doc1"], metadatas=["test"])
#132
#133	assert app1.db.count() == 1
#134	assert app2.db.count() == 0
#135
#136	app1.db.add(ids=["1", "2"], documents=["doc1", "doc2"], metadatas=["test", "test"])
#137	app2.db.add(ids=["0"], documents=["doc1"], metadatas=["test"])
#138
#139	app1.set_collection_name("test_collection_2")
#140	assert app1.db.count() == 1
#141	app2.set_collection_name("test_collection_1")
#142	assert app2.db.count() == 3
#143
#144	# cleanup
#145	app1.db.reset()
#146	app2.db.reset()
#147
#148
#149	def test_lancedb_collection_ids_share_collections():
#150	db1 = LanceDB(config=LanceDBConfig(allow_reset=True, dir="test-db"))
#151	app1 = App(config=AppConfig(collect_metrics=False), db=db1)
#152	app1.set_collection_name("one_collection")
#153	db2 = LanceDB(config=LanceDBConfig(allow_reset=True, dir="test-db"))
#154	app2 = App(config=AppConfig(collect_metrics=False), db=db2)
#155	app2.set_collection_name("one_collection")
#156
#157	# cleanup
#158	app1.db.reset()
#159	app2.db.reset()
#160
#161	app1.db.add(ids=["0", "1"], documents=["doc1", "doc2"], metadatas=["test", "test"])
#162	app2.db.add(ids=["2"], documents=["doc3"], metadatas=["test"])
#163
#164	assert app1.db.count() == 2
#165	assert app2.db.count() == 3
#166
#167	# cleanup
#168	app1.db.reset()
#169	app2.db.reset()
#170
#171
#172	def test_lancedb_collection_reset():
#173	db1 = LanceDB(config=LanceDBConfig(allow_reset=True, dir="test-db"))
#174	app1 = App(config=AppConfig(collect_metrics=False), db=db1)
#175	app1.set_collection_name("one_collection")
#176	db2 = LanceDB(config=LanceDBConfig(allow_reset=True, dir="test-db"))
#177	app2 = App(config=AppConfig(collect_metrics=False), db=db2)
#178	app2.set_collection_name("two_collection")
#179	db3 = LanceDB(config=LanceDBConfig(allow_reset=True, dir="test-db"))
#180	app3 = App(config=AppConfig(collect_metrics=False), db=db3)
#181	app3.set_collection_name("three_collection")
#182	db4 = LanceDB(config=LanceDBConfig(allow_reset=True, dir="test-db"))
#183	app4 = App(config=AppConfig(collect_metrics=False), db=db4)
#184	app4.set_collection_name("four_collection")
#185
#186	# cleanup if any previous tests failed or were interrupted
#187	app1.db.reset()
#188	app2.db.reset()
#189	app3.db.reset()
#190	app4.db.reset()
#191
#192	app1.db.add(ids=["1"], documents=["doc1"], metadatas=["test"])
#193	app2.db.add(ids=["2"], documents=["doc2"], metadatas=["test"])
#194	app3.db.add(ids=["3"], documents=["doc3"], metadatas=["test"])
#195	app4.db.add(ids=["4"], documents=["doc4"], metadatas=["test"])
#196
#197	app1.db.reset()
#198
#199	assert app1.db.count() == 0
#200	assert app2.db.count() == 1
#201	assert app3.db.count() == 1
#202	assert app4.db.count() == 1
#203
#204	# cleanup
#205	app2.db.reset()
#206	app3.db.reset()
#207	app4.db.reset()
#208
#209
#210	def generate_embeddings(dummy_embed, embed_size):
#211	generated_embedding = []
#212	for i in range(embed_size):
#213	generated_embedding.append(dummy_embed)
#214
#215	return generated_embedding
#216

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public