增加单元测试

hhyo · Sep 2, 2024 · e9f374a · e9f374a
1 parent 64e36ff
commit e9f374a
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 1 deletion.
diff --git a/sql/engines/elasticsearch.py b/sql/engines/elasticsearch.py
@@ -143,8 +143,8 @@ def get_all_tables(self, db_name, **kwargs):
                     # 将系统的API作为表名
                     tables.add("/_cat/indices/" + self.db_name)
                     tables.add("/_cat/nodes")
-                    tables.add("/_security/user")
                     tables.add("/_security/role")
+                    tables.add("/_security/user")
 
                 for index_name in indices.keys():
                     if index_name.startswith("."):

diff --git a/sql/engines/test_elasticsearch.py b/sql/engines/test_elasticsearch.py
@@ -55,6 +55,32 @@ def test_get_all_tables(self, mockElasticsearch):
         result = self.engine.get_all_tables(db_name="test")
         self.assertEqual(result.rows, ["test__index1", "test__index2"])
 
+    @patch("sql.engines.elasticsearch.Elasticsearch")
+    def test_get_all_tables_system(self, mockElasticsearch):
+        """测试获取所有表名，特定数据库 'system'"""
+        mock_conn = Mock()
+        mockElasticsearch.return_value = mock_conn
+
+        # 假设系统相关的索引（以.开头的）和特定的_cat API端点
+        mock_conn.indices.get_alias.return_value = {
+            ".kibana_1": {},
+            ".security": {},
+            "test__index": {},
+        }
+
+        result = self.engine.get_all_tables(db_name="system")
+
+        # 预期结果应包括系统相关的表名和 /_cat API 端点
+        expected_tables = [
+            "/_cat/indices/*",
+            "/_cat/indices/test",
+            "/_cat/nodes",
+            "/_security/role",
+            "/_security/user",
+        ]
+
+        self.assertEqual(result.rows, expected_tables)
+
     @patch("sql.engines.elasticsearch.Elasticsearch")
     def test_query(self, mockElasticsearch):
         mock_conn = Mock()