Skip to content

Commit ac695b0

Browse files
Zhi-aZhi-a
committed
✨Added Datamate vector knowledge base core v2
1 parent f5f473b commit ac695b0

File tree

1 file changed

+30
-23
lines changed

1 file changed

+30
-23
lines changed

‎test/backend/services/test_vectordatabase_service.py‎

Lines changed: 30 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -404,8 +404,9 @@ def test_list_indices_without_stats(self, mock_get_knowledge, mock_get_user_tena
404404
self.mock_vdb_core.get_user_indices.return_value = ["index1", "index2"]
405405
mock_get_knowledge.return_value = [
406406
{"index_name": "index1",
407-
"embedding_model_name": "test-model", "group_ids": "1,2"},
408-
{"index_name": "index2", "embedding_model_name": "test-model", "group_ids": ""}
407+
"embedding_model_name": "test-model", "group_ids": "1,2", "knowledge_sources": "elasticsearch"},
408+
{"index_name": "index2", "embedding_model_name": "test-model",
409+
"group_ids": "", "knowledge_sources": "elasticsearch"}
409410
]
410411
mock_get_user_tenant.return_value = {
411412
"user_role": "SU", "tenant_id": "test_tenant"}
@@ -446,8 +447,9 @@ def test_list_indices_with_stats(self, mock_get_knowledge, mock_get_user_tenant,
446447
}
447448
mock_get_knowledge.return_value = [
448449
{"index_name": "index1",
449-
"embedding_model_name": "test-model", "group_ids": "1,2"},
450-
{"index_name": "index2", "embedding_model_name": "test-model", "group_ids": ""}
450+
"embedding_model_name": "test-model", "group_ids": "1,2", "knowledge_sources": "elasticsearch"},
451+
{"index_name": "index2", "embedding_model_name": "test-model",
452+
"group_ids": "", "knowledge_sources": "elasticsearch"}
451453
]
452454
mock_get_user_tenant.return_value = {
453455
"user_role": "SU", "tenant_id": "test_tenant"}
@@ -486,7 +488,7 @@ def test_list_indices_skips_missing_indices(self, mock_get_info, mock_get_user_t
486488
self.mock_vdb_core.get_user_indices.return_value = ["es_index"]
487489
mock_get_info.return_value = [
488490
{"index_name": "dangling_index",
489-
"embedding_model_name": "model-A", "group_ids": "1"}
491+
"embedding_model_name": "model-A", "group_ids": "1", "knowledge_sources": "elasticsearch"}
490492
]
491493
mock_get_user_tenant.return_value = {
492494
"user_role": "SU", "tenant_id": "tenant-1"}
@@ -513,7 +515,8 @@ def test_list_indices_stats_defaults_when_missing(self, mock_get_info, mock_get_
513515
"""
514516
self.mock_vdb_core.get_user_indices.return_value = ["index1"]
515517
mock_get_info.return_value = [
516-
{"index_name": "index1", "embedding_model_name": "model-A", "group_ids": "1,2"}
518+
{"index_name": "index1", "embedding_model_name": "model-A",
519+
"group_ids": "1,2", "knowledge_sources": "elasticsearch"}
517520
]
518521
self.mock_vdb_core.get_indices_detail.return_value = {}
519522
mock_get_user_tenant.return_value = {
@@ -542,7 +545,8 @@ def test_list_indices_backfills_missing_model_names(self, mock_get_info, mock_up
542545
"""
543546
self.mock_vdb_core.get_user_indices.return_value = ["index1"]
544547
mock_get_info.return_value = [
545-
{"index_name": "index1", "embedding_model_name": None}
548+
{"index_name": "index1", "embedding_model_name": None,
549+
"knowledge_sources": "elasticsearch"}
546550
]
547551
self.mock_vdb_core.get_indices_detail.return_value = {
548552
"index1": {"base_info": {"embedding_model": "text-embedding-ada-002"}}
@@ -574,7 +578,8 @@ def test_list_indices_stats_surfaces_elasticsearch_errors(self, mock_get_info, m
574578
"""
575579
self.mock_vdb_core.get_user_indices.return_value = ["index1"]
576580
mock_get_info.return_value = [
577-
{"index_name": "index1", "embedding_model_name": "model-A", "group_ids": "1,2"}
581+
{"index_name": "index1", "embedding_model_name": "model-A",
582+
"group_ids": "1,2", "knowledge_sources": "elasticsearch"}
578583
]
579584
self.mock_vdb_core.get_indices_detail.side_effect = Exception(
580585
"503 Service Unavailable"
@@ -603,7 +608,8 @@ def test_list_indices_stats_keeps_non_stat_fields(self, mock_get_info, mock_get_
603608
"""
604609
self.mock_vdb_core.get_user_indices.return_value = ["index1"]
605610
mock_get_info.return_value = [
606-
{"index_name": "index1", "embedding_model_name": "model-A", "group_ids": "1,2"}
611+
{"index_name": "index1", "embedding_model_name": "model-A",
612+
"group_ids": "1,2", "knowledge_sources": "elasticsearch"}
607613
]
608614
detailed_stats = {
609615
"index1": {
@@ -652,15 +658,17 @@ def test_list_indices_creator_permission(self, mock_get_knowledge, mock_get_user
652658
"group_ids": "1",
653659
"created_by": "test_user", # User is creator
654660
"ingroup_permission": "READ_ONLY",
655-
"tenant_id": "test_tenant"
661+
"tenant_id": "test_tenant",
662+
"knowledge_sources": "elasticsearch"
656663
},
657664
{
658665
"index_name": "index2",
659666
"embedding_model_name": "test-model",
660667
"group_ids": "1",
661668
"created_by": "other_user", # User is not creator
662669
"ingroup_permission": "EDIT",
663-
"tenant_id": "test_tenant"
670+
"tenant_id": "test_tenant",
671+
"knowledge_sources": "elasticsearch"
664672
}
665673
]
666674
mock_get_user_tenant.return_value = {
@@ -704,13 +712,15 @@ def test_list_indices_fallback_admin_logic(self, mock_get_knowledge, mock_get_us
704712
"index_name": "index1",
705713
"embedding_model_name": "test-model",
706714
"group_ids": "1,2",
707-
"tenant_id": "legacy_admin_user" # Same as user_id
715+
"tenant_id": "legacy_admin_user", # Same as user_id
716+
"knowledge_sources": "elasticsearch"
708717
},
709718
{
710719
"index_name": "index2",
711720
"embedding_model_name": "test-model",
712721
"group_ids": "3",
713-
"tenant_id": "legacy_admin_user" # Same as user_id
722+
"tenant_id": "legacy_admin_user", # Same as user_id
723+
"knowledge_sources": "elasticsearch"
714724
}
715725
]
716726
# user_role is None to test fallback logic
@@ -762,13 +772,15 @@ def test_list_indices_speed_version_admin_logic(self, mock_get_knowledge, mock_g
762772
"index_name": "index1",
763773
"embedding_model_name": "test-model",
764774
"group_ids": "1,2",
765-
"tenant_id": "tenant_id" # DEFAULT_TENANT_ID
775+
"tenant_id": "tenant_id", # DEFAULT_TENANT_ID
776+
"knowledge_sources": "elasticsearch"
766777
},
767778
{
768779
"index_name": "index2",
769780
"embedding_model_name": "test-model",
770781
"group_ids": "3",
771-
"tenant_id": "tenant_id" # DEFAULT_TENANT_ID
782+
"tenant_id": "tenant_id", # DEFAULT_TENANT_ID
783+
"knowledge_sources": "elasticsearch"
772784
}
773785
]
774786
# user_role is USER but should be overridden by SPEED logic
@@ -2243,8 +2255,9 @@ def test_list_indices_success_status_200(self, mock_response, mock_get_knowledge
22432255
mock_response.status_code = 200
22442256
mock_get_knowledge.return_value = [
22452257
{"index_name": "index1",
2246-
"embedding_model_name": "test-model", "group_ids": "1,2"},
2247-
{"index_name": "index2", "embedding_model_name": "test-model", "group_ids": ""}
2258+
"embedding_model_name": "test-model", "group_ids": "1,2", "knowledge_sources": "elasticsearch"},
2259+
{"index_name": "index2", "embedding_model_name": "test-model",
2260+
"group_ids": "", "knowledge_sources": "elasticsearch"}
22482261
]
22492262
mock_get_user_tenant.return_value = {
22502263
"user_role": "SU", "tenant_id": "test_tenant"}
@@ -2525,12 +2538,6 @@ def test_check_kb_exist_exists_in_tenant(self, mock_get_knowledge):
25252538
})
25262539
self.assertEqual(result["status"], "exists_in_tenant")
25272540

2528-
2529-
2530-
2531-
2532-
2533-
25342541
# Note: generate_knowledge_summary_stream function has been removed
25352542
# These tests are no longer relevant as the function was replaced with summary_index_name
25362543

0 commit comments

Comments
 (0)