Skip to content

Commit f1be313

Browse files
committed
repair test case improve
1 parent aa3b7d0 commit f1be313

File tree

2 files changed

+533
-1
lines changed

2 files changed

+533
-1
lines changed

test/backend/services/test_elasticsearch_service.py

Lines changed: 218 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1778,6 +1778,224 @@ def test_check_kb_exist_orphan_in_pg_delete_failure(self, mock_get_knowledge, mo
17781778
# Note: generate_knowledge_summary_stream function has been removed
17791779
# These tests are no longer relevant as the function was replaced with summary_index_name
17801780

1781+
def test_get_es_core(self):
1782+
"""
1783+
Test get_es_core function returns the elastic_core instance.
1784+
1785+
This test verifies that:
1786+
1. The get_es_core function returns the correct elastic_core instance
1787+
2. The function is properly imported and accessible
1788+
"""
1789+
from backend.services.elasticsearch_service import get_es_core
1790+
1791+
# Execute
1792+
result = get_es_core()
1793+
1794+
# Assert
1795+
self.assertIsNotNone(result)
1796+
# The result should be the elastic_core instance
1797+
self.assertTrue(hasattr(result, 'client'))
1798+
1799+
@patch('backend.services.elasticsearch_service.tenant_config_manager')
1800+
def test_get_embedding_model_embedding_type(self, mock_tenant_config_manager):
1801+
"""
1802+
Test get_embedding_model with embedding model type.
1803+
1804+
This test verifies that:
1805+
1. When model_type is "embedding", OpenAICompatibleEmbedding is returned
1806+
2. The correct parameters are passed to the embedding model
1807+
"""
1808+
# Setup
1809+
mock_config = {
1810+
"model_type": "embedding",
1811+
"api_key": "test_api_key",
1812+
"base_url": "https://test.api.com",
1813+
"model_name": "test-model",
1814+
"max_tokens": 1024
1815+
}
1816+
mock_tenant_config_manager.get_model_config.return_value = mock_config
1817+
1818+
# Stop the mock from setUp to test the real function
1819+
self.get_embedding_model_patcher.stop()
1820+
1821+
try:
1822+
with patch('backend.services.elasticsearch_service.OpenAICompatibleEmbedding') as mock_embedding_class:
1823+
mock_embedding_instance = MagicMock()
1824+
mock_embedding_class.return_value = mock_embedding_instance
1825+
1826+
with patch('backend.services.elasticsearch_service.get_model_name_from_config') as mock_get_model_name:
1827+
mock_get_model_name.return_value = "test-model"
1828+
1829+
# Execute - now we can call the real function
1830+
from backend.services.elasticsearch_service import get_embedding_model
1831+
result = get_embedding_model("test_tenant")
1832+
1833+
# Assert
1834+
self.assertEqual(result, mock_embedding_instance)
1835+
mock_tenant_config_manager.get_model_config.assert_called_once_with(
1836+
key="EMBEDDING_ID", tenant_id="test_tenant")
1837+
mock_embedding_class.assert_called_once_with(
1838+
api_key="test_api_key",
1839+
base_url="https://test.api.com",
1840+
model_name="test-model",
1841+
embedding_dim=1024
1842+
)
1843+
finally:
1844+
# Restart the mock for other tests
1845+
self.get_embedding_model_patcher.start()
1846+
1847+
@patch('backend.services.elasticsearch_service.tenant_config_manager')
1848+
def test_get_embedding_model_multi_embedding_type(self, mock_tenant_config_manager):
1849+
"""
1850+
Test get_embedding_model with multi_embedding model type.
1851+
1852+
This test verifies that:
1853+
1. When model_type is "multi_embedding", JinaEmbedding is returned
1854+
2. The correct parameters are passed to the embedding model
1855+
"""
1856+
# Setup
1857+
mock_config = {
1858+
"model_type": "multi_embedding",
1859+
"api_key": "test_api_key",
1860+
"base_url": "https://test.api.com",
1861+
"model_name": "test-model",
1862+
"max_tokens": 2048
1863+
}
1864+
mock_tenant_config_manager.get_model_config.return_value = mock_config
1865+
1866+
# Stop the mock from setUp to test the real function
1867+
self.get_embedding_model_patcher.stop()
1868+
1869+
try:
1870+
with patch('backend.services.elasticsearch_service.JinaEmbedding') as mock_embedding_class:
1871+
mock_embedding_instance = MagicMock()
1872+
mock_embedding_class.return_value = mock_embedding_instance
1873+
1874+
with patch('backend.services.elasticsearch_service.get_model_name_from_config') as mock_get_model_name:
1875+
mock_get_model_name.return_value = "test-model"
1876+
1877+
# Execute - now we can call the real function
1878+
from backend.services.elasticsearch_service import get_embedding_model
1879+
result = get_embedding_model("test_tenant")
1880+
1881+
# Assert
1882+
self.assertEqual(result, mock_embedding_instance)
1883+
mock_tenant_config_manager.get_model_config.assert_called_once_with(
1884+
key="EMBEDDING_ID", tenant_id="test_tenant")
1885+
mock_embedding_class.assert_called_once_with(
1886+
api_key="test_api_key",
1887+
base_url="https://test.api.com",
1888+
model_name="test-model",
1889+
embedding_dim=2048
1890+
)
1891+
finally:
1892+
# Restart the mock for other tests
1893+
self.get_embedding_model_patcher.start()
1894+
1895+
@patch('backend.services.elasticsearch_service.tenant_config_manager')
1896+
def test_get_embedding_model_unknown_type(self, mock_tenant_config_manager):
1897+
"""
1898+
Test get_embedding_model with unknown model type.
1899+
1900+
This test verifies that:
1901+
1. When model_type is neither "embedding" nor "multi_embedding", None is returned
1902+
2. The function handles unknown model types gracefully
1903+
"""
1904+
# Setup
1905+
mock_config = {
1906+
"model_type": "unknown_type",
1907+
"api_key": "test_api_key",
1908+
"base_url": "https://test.api.com",
1909+
"model_name": "test-model",
1910+
"max_tokens": 1024
1911+
}
1912+
mock_tenant_config_manager.get_model_config.return_value = mock_config
1913+
1914+
# Stop the mock from setUp to test the real function
1915+
self.get_embedding_model_patcher.stop()
1916+
1917+
try:
1918+
# Execute - now we can call the real function
1919+
from backend.services.elasticsearch_service import get_embedding_model
1920+
result = get_embedding_model("test_tenant")
1921+
1922+
# Assert
1923+
self.assertIsNone(result)
1924+
mock_tenant_config_manager.get_model_config.assert_called_once_with(
1925+
key="EMBEDDING_ID", tenant_id="test_tenant")
1926+
finally:
1927+
# Restart the mock for other tests
1928+
self.get_embedding_model_patcher.start()
1929+
1930+
@patch('backend.services.elasticsearch_service.tenant_config_manager')
1931+
def test_get_embedding_model_empty_type(self, mock_tenant_config_manager):
1932+
"""
1933+
Test get_embedding_model with empty model type.
1934+
1935+
This test verifies that:
1936+
1. When model_type is empty string, None is returned
1937+
2. The function handles empty model types gracefully
1938+
"""
1939+
# Setup
1940+
mock_config = {
1941+
"model_type": "",
1942+
"api_key": "test_api_key",
1943+
"base_url": "https://test.api.com",
1944+
"model_name": "test-model",
1945+
"max_tokens": 1024
1946+
}
1947+
mock_tenant_config_manager.get_model_config.return_value = mock_config
1948+
1949+
# Stop the mock from setUp to test the real function
1950+
self.get_embedding_model_patcher.stop()
1951+
1952+
try:
1953+
# Execute - now we can call the real function
1954+
from backend.services.elasticsearch_service import get_embedding_model
1955+
result = get_embedding_model("test_tenant")
1956+
1957+
# Assert
1958+
self.assertIsNone(result)
1959+
mock_tenant_config_manager.get_model_config.assert_called_once_with(
1960+
key="EMBEDDING_ID", tenant_id="test_tenant")
1961+
finally:
1962+
# Restart the mock for other tests
1963+
self.get_embedding_model_patcher.start()
1964+
1965+
@patch('backend.services.elasticsearch_service.tenant_config_manager')
1966+
def test_get_embedding_model_missing_type(self, mock_tenant_config_manager):
1967+
"""
1968+
Test get_embedding_model with missing model type.
1969+
1970+
This test verifies that:
1971+
1. When model_type is missing from config, None is returned
1972+
2. The function handles missing model types gracefully
1973+
"""
1974+
# Setup
1975+
mock_config = {
1976+
"api_key": "test_api_key",
1977+
"base_url": "https://test.api.com",
1978+
"model_name": "test-model",
1979+
"max_tokens": 1024
1980+
}
1981+
mock_tenant_config_manager.get_model_config.return_value = mock_config
1982+
1983+
# Stop the mock from setUp to test the real function
1984+
self.get_embedding_model_patcher.stop()
1985+
1986+
try:
1987+
# Execute - now we can call the real function
1988+
from backend.services.elasticsearch_service import get_embedding_model
1989+
result = get_embedding_model("test_tenant")
1990+
1991+
# Assert
1992+
self.assertIsNone(result)
1993+
mock_tenant_config_manager.get_model_config.assert_called_once_with(
1994+
key="EMBEDDING_ID", tenant_id="test_tenant")
1995+
finally:
1996+
# Restart the mock for other tests
1997+
self.get_embedding_model_patcher.start()
1998+
17811999

17822000
if __name__ == '__main__':
17832001
unittest.main()

0 commit comments

Comments
 (0)