my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

my-project-public — gitlawb

#1	import os
#2	import shutil
#3	from unittest.mock import patch
#4
#5	import pytest
#6	from chromadb.config import Settings
#7
#8	from embedchain import App
#9	from embedchain.config import AppConfig, ChromaDbConfig
#10	from embedchain.vectordb.chroma import ChromaDB
#11
#12	os.environ["OPENAI_API_KEY"] = "test-api-key"
#13
#14
#15	@pytest.fixture
#16	def chroma_db():
#17	return ChromaDB(config=ChromaDbConfig(host="test-host", port="1234"))
#18
#19
#20	@pytest.fixture
#21	def app_with_settings():
#22	chroma_config = ChromaDbConfig(allow_reset=True, dir="test-db")
#23	chroma_db = ChromaDB(config=chroma_config)
#24	app_config = AppConfig(collect_metrics=False)
#25	return App(config=app_config, db=chroma_db)
#26
#27
#28	@pytest.fixture(scope="session", autouse=True)
#29	def cleanup_db():
#30	yield
#31	try:
#32	shutil.rmtree("test-db")
#33	except OSError as e:
#34	print("Error: %s - %s." % (e.filename, e.strerror))
#35
#36
#37	@patch("embedchain.vectordb.chroma.chromadb.Client")
#38	def test_chroma_db_init_with_host_and_port(mock_client):
#39	chroma_db = ChromaDB(config=ChromaDbConfig(host="test-host", port="1234")) # noqa
#40	called_settings: Settings = mock_client.call_args[0][0]
#41	assert called_settings.chroma_server_host == "test-host"
#42	assert called_settings.chroma_server_http_port == "1234"
#43
#44
#45	@patch("embedchain.vectordb.chroma.chromadb.Client")
#46	def test_chroma_db_init_with_basic_auth(mock_client):
#47	chroma_config = {
#48	"host": "test-host",
#49	"port": "1234",
#50	"chroma_settings": {
#51	"chroma_client_auth_provider": "chromadb.auth.basic.BasicAuthClientProvider",
#52	"chroma_client_auth_credentials": "admin:admin",
#53	},
#54	}
#55
#56	ChromaDB(config=ChromaDbConfig(**chroma_config))
#57	called_settings: Settings = mock_client.call_args[0][0]
#58	assert called_settings.chroma_server_host == "test-host"
#59	assert called_settings.chroma_server_http_port == "1234"
#60	assert (
#61	called_settings.chroma_client_auth_provider == chroma_config["chroma_settings"]["chroma_client_auth_provider"]
#62	)
#63	assert (
#64	called_settings.chroma_client_auth_credentials
#65	== chroma_config["chroma_settings"]["chroma_client_auth_credentials"]
#66	)
#67
#68
#69	@patch("embedchain.vectordb.chroma.chromadb.Client")
#70	def test_app_init_with_host_and_port(mock_client):
#71	host = "test-host"
#72	port = "1234"
#73	config = AppConfig(collect_metrics=False)
#74	db_config = ChromaDbConfig(host=host, port=port)
#75	db = ChromaDB(config=db_config)
#76	_app = App(config=config, db=db)
#77
#78	called_settings: Settings = mock_client.call_args[0][0]
#79	assert called_settings.chroma_server_host == host
#80	assert called_settings.chroma_server_http_port == port
#81
#82
#83	@patch("embedchain.vectordb.chroma.chromadb.Client")
#84	def test_app_init_with_host_and_port_none(mock_client):
#85	db = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db"))
#86	_app = App(config=AppConfig(collect_metrics=False), db=db)
#87
#88	called_settings: Settings = mock_client.call_args[0][0]
#89	assert called_settings.chroma_server_host is None
#90	assert called_settings.chroma_server_http_port is None
#91
#92
#93	def test_chroma_db_duplicates_throw_warning(caplog):
#94	db = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db"))
#95	app = App(config=AppConfig(collect_metrics=False), db=db)
#96	app.db.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
#97	app.db.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
#98	assert "Insert of existing embedding ID: 0" in caplog.text
#99	assert "Add of existing embedding ID: 0" in caplog.text
#100	app.db.reset()
#101
#102
#103	def test_chroma_db_duplicates_collections_no_warning(caplog):
#104	db = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db"))
#105	app = App(config=AppConfig(collect_metrics=False), db=db)
#106	app.set_collection_name("test_collection_1")
#107	app.db.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
#108	app.set_collection_name("test_collection_2")
#109	app.db.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
#110	assert "Insert of existing embedding ID: 0" not in caplog.text
#111	assert "Add of existing embedding ID: 0" not in caplog.text
#112	app.db.reset()
#113	app.set_collection_name("test_collection_1")
#114	app.db.reset()
#115
#116
#117	def test_chroma_db_collection_init_with_default_collection():
#118	db = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db"))
#119	app = App(config=AppConfig(collect_metrics=False), db=db)
#120	assert app.db.collection.name == "embedchain_store"
#121
#122
#123	def test_chroma_db_collection_init_with_custom_collection():
#124	db = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db"))
#125	app = App(config=AppConfig(collect_metrics=False), db=db)
#126	app.set_collection_name(name="test_collection")
#127	assert app.db.collection.name == "test_collection"
#128
#129
#130	def test_chroma_db_collection_set_collection_name():
#131	db = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db"))
#132	app = App(config=AppConfig(collect_metrics=False), db=db)
#133	app.set_collection_name("test_collection")
#134	assert app.db.collection.name == "test_collection"
#135
#136
#137	def test_chroma_db_collection_changes_encapsulated():
#138	db = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db"))
#139	app = App(config=AppConfig(collect_metrics=False), db=db)
#140	app.set_collection_name("test_collection_1")
#141	assert app.db.count() == 0
#142
#143	app.db.collection.add(embeddings=[0, 0, 0], ids=["0"])
#144	assert app.db.count() == 1
#145
#146	app.set_collection_name("test_collection_2")
#147	assert app.db.count() == 0
#148
#149	app.db.collection.add(embeddings=[0, 0, 0], ids=["0"])
#150	app.set_collection_name("test_collection_1")
#151	assert app.db.count() == 1
#152	app.db.reset()
#153	app.set_collection_name("test_collection_2")
#154	app.db.reset()
#155
#156
#157	def test_chroma_db_collection_collections_are_persistent():
#158	db = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db"))
#159	app = App(config=AppConfig(collect_metrics=False), db=db)
#160	app.set_collection_name("test_collection_1")
#161	app.db.collection.add(embeddings=[[0, 0, 0]], ids=["0"])
#162	del app
#163
#164	db = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db"))
#165	app = App(config=AppConfig(collect_metrics=False), db=db)
#166	app.set_collection_name("test_collection_1")
#167	assert app.db.count() == 1
#168
#169	app.db.reset()
#170
#171
#172	def test_chroma_db_collection_parallel_collections():
#173	db1 = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db", collection_name="test_collection_1"))
#174	app1 = App(
#175	config=AppConfig(collect_metrics=False),
#176	db=db1,
#177	)
#178	db2 = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db", collection_name="test_collection_2"))
#179	app2 = App(
#180	config=AppConfig(collect_metrics=False),
#181	db=db2,
#182	)
#183
#184	# cleanup if any previous tests failed or were interrupted
#185	app1.db.reset()
#186	app2.db.reset()
#187
#188	app1.db.collection.add(embeddings=[0, 0, 0], ids=["0"])
#189	assert app1.db.count() == 1
#190	assert app2.db.count() == 0
#191
#192	app1.db.collection.add(embeddings=[[0, 0, 0], [1, 1, 1]], ids=["1", "2"])
#193	app2.db.collection.add(embeddings=[0, 0, 0], ids=["0"])
#194
#195	app1.set_collection_name("test_collection_2")
#196	assert app1.db.count() == 1
#197	app2.set_collection_name("test_collection_1")
#198	assert app2.db.count() == 3
#199
#200	# cleanup
#201	app1.db.reset()
#202	app2.db.reset()
#203
#204
#205	def test_chroma_db_collection_ids_share_collections():
#206	db1 = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db"))
#207	app1 = App(config=AppConfig(collect_metrics=False), db=db1)
#208	app1.set_collection_name("one_collection")
#209	db2 = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db"))
#210	app2 = App(config=AppConfig(collect_metrics=False), db=db2)
#211	app2.set_collection_name("one_collection")
#212
#213	app1.db.collection.add(embeddings=[[0, 0, 0], [1, 1, 1]], ids=["0", "1"])
#214	app2.db.collection.add(embeddings=[0, 0, 0], ids=["2"])
#215
#216	assert app1.db.count() == 3
#217	assert app2.db.count() == 3
#218
#219	# cleanup
#220	app1.db.reset()
#221	app2.db.reset()
#222
#223
#224	def test_chroma_db_collection_reset():
#225	db1 = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db"))
#226	app1 = App(config=AppConfig(collect_metrics=False), db=db1)
#227	app1.set_collection_name("one_collection")
#228	db2 = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db"))
#229	app2 = App(config=AppConfig(collect_metrics=False), db=db2)
#230	app2.set_collection_name("two_collection")
#231	db3 = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db"))
#232	app3 = App(config=AppConfig(collect_metrics=False), db=db3)
#233	app3.set_collection_name("three_collection")
#234	db4 = ChromaDB(config=ChromaDbConfig(allow_reset=True, dir="test-db"))
#235	app4 = App(config=AppConfig(collect_metrics=False), db=db4)
#236	app4.set_collection_name("four_collection")
#237
#238	app1.db.collection.add(embeddings=[0, 0, 0], ids=["1"])
#239	app2.db.collection.add(embeddings=[0, 0, 0], ids=["2"])
#240	app3.db.collection.add(embeddings=[0, 0, 0], ids=["3"])
#241	app4.db.collection.add(embeddings=[0, 0, 0], ids=["4"])
#242
#243	app1.db.reset()
#244
#245	assert app1.db.count() == 0
#246	assert app2.db.count() == 1
#247	assert app3.db.count() == 1
#248	assert app4.db.count() == 1
#249
#250	# cleanup
#251	app2.db.reset()
#252	app3.db.reset()
#253	app4.db.reset()
#254

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public