73 lines
2.3 KiB
Python
73 lines
2.3 KiB
Python
# ruff: noqa: E501
|
|
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from embedchain import App
|
|
from embedchain.config import InitConfig
|
|
from embedchain.vectordb.chroma_db import ChromaDB, chromadb
|
|
|
|
|
|
class TestChromaDbHosts(unittest.TestCase):
|
|
def test_init_with_host_and_port(self):
|
|
"""
|
|
Test if the `ChromaDB` instance is initialized with the correct host and port values.
|
|
"""
|
|
host = "test-host"
|
|
port = "1234"
|
|
|
|
with patch.object(chromadb, "Client") as mock_client:
|
|
_db = ChromaDB(host=host, port=port)
|
|
|
|
expected_settings = chromadb.config.Settings(
|
|
chroma_api_impl="rest",
|
|
chroma_server_host=host,
|
|
chroma_server_http_port=port,
|
|
)
|
|
|
|
mock_client.assert_called_once_with(expected_settings)
|
|
|
|
|
|
class TestChromaDbHostsInit(unittest.TestCase):
|
|
@patch("embedchain.vectordb.chroma_db.chromadb.Client")
|
|
def test_init_with_host_and_port(self, mock_client):
|
|
"""
|
|
Test if the `App` instance is initialized with the correct host and port values.
|
|
"""
|
|
host = "test-host"
|
|
port = "1234"
|
|
|
|
config = InitConfig(host=host, port=port)
|
|
|
|
_app = App(config)
|
|
|
|
self.assertEqual(mock_client.call_args[0][0].chroma_server_host, host)
|
|
self.assertEqual(mock_client.call_args[0][0].chroma_server_http_port, port)
|
|
|
|
|
|
class TestChromaDbHostsNone(unittest.TestCase):
|
|
@patch("embedchain.vectordb.chroma_db.chromadb.Client")
|
|
def test_init_with_host_and_port(self, mock_client):
|
|
"""
|
|
Test if the `App` instance is initialized without default hosts and ports.
|
|
"""
|
|
|
|
_app = App()
|
|
|
|
self.assertEqual(mock_client.call_args[0][0].chroma_server_host, None)
|
|
self.assertEqual(mock_client.call_args[0][0].chroma_server_http_port, None)
|
|
|
|
|
|
class TestChromaDbHostsLoglevel(unittest.TestCase):
|
|
@patch("embedchain.vectordb.chroma_db.chromadb.Client")
|
|
def test_init_with_host_and_port(self, mock_client):
|
|
"""
|
|
Test if the `App` instance is initialized without a config that does not contain default hosts and ports.
|
|
"""
|
|
config = InitConfig(log_level="DEBUG")
|
|
|
|
_app = App(config)
|
|
|
|
self.assertEqual(mock_client.call_args[0][0].chroma_server_host, None)
|
|
self.assertEqual(mock_client.call_args[0][0].chroma_server_http_port, None)
|