From 5fc11fb70681d71271af76fa374fbafebf3744a0 Mon Sep 17 00:00:00 2001
From: Vitor Avila <96086495+Vitor-Avila@users.noreply.github.com>
Date: Fri, 31 Jan 2025 08:36:09 -0300
Subject: [PATCH] chore: Add more database-related tests (follow up to #31948)
 (#32054)

---
 superset/commands/database/tables.py          |  6 ++
 .../commands/databases/tables_test.py         | 28 +++++-----
 tests/unit_tests/models/core_test.py          | 56 +++++++++++++++++++
 3 files changed, 76 insertions(+), 14 deletions(-)

diff --git a/superset/commands/database/tables.py b/superset/commands/database/tables.py
index c8fa88400..c0c0507eb 100644
--- a/superset/commands/database/tables.py
+++ b/superset/commands/database/tables.py
@@ -60,6 +60,9 @@ class TablesDatabaseCommand(BaseCommand):
                 catalog=self._catalog_name,
                 schema=self._schema_name,
                 datasource_names=sorted(
+                    # get_all_table_names_in_schema may return raw (unserialized) cached
+                    # results, so we wrap them as DatasourceName objects here instead of
+                    # directly in the method to ensure consistency.
                     DatasourceName(*datasource_name)
                     for datasource_name in self._model.get_all_table_names_in_schema(
                         catalog=self._catalog_name,
@@ -76,6 +79,9 @@ class TablesDatabaseCommand(BaseCommand):
                 catalog=self._catalog_name,
                 schema=self._schema_name,
                 datasource_names=sorted(
+                    # get_all_view_names_in_schema may return raw (unserialized) cached
+                    # results, so we wrap them as DatasourceName objects here instead of
+                    # directly in the method to ensure consistency.
                     DatasourceName(*datasource_name)
                     for datasource_name in self._model.get_all_view_names_in_schema(
                         catalog=self._catalog_name,
diff --git a/tests/unit_tests/commands/databases/tables_test.py b/tests/unit_tests/commands/databases/tables_test.py
index db446b46b..d7eaecfa5 100644
--- a/tests/unit_tests/commands/databases/tables_test.py
+++ b/tests/unit_tests/commands/databases/tables_test.py
@@ -34,13 +34,13 @@ def database_with_catalog(mocker: MockerFixture) -> MagicMock:
 
     database = mocker.MagicMock()
     database.database_name = "test_database"
-    database.get_all_table_names_in_schema.return_value = [
-        DatasourceName("table1", "schema1", "catalog1"),
-        DatasourceName("table2", "schema1", "catalog1"),
-    ]
-    database.get_all_view_names_in_schema.return_value = [
-        DatasourceName("view1", "schema1", "catalog1"),
-    ]
+    database.get_all_table_names_in_schema.return_value = {
+        ("table1", "schema1", "catalog1"),
+        ("table2", "schema1", "catalog1"),
+    }
+    database.get_all_view_names_in_schema.return_value = {
+        ("view1", "schema1", "catalog1"),
+    }
 
     DatabaseDAO = mocker.patch("superset.commands.database.tables.DatabaseDAO")  # noqa: N806
     DatabaseDAO.find_by_id.return_value = database
@@ -57,13 +57,13 @@ def database_without_catalog(mocker: MockerFixture) -> MagicMock:
 
     database = mocker.MagicMock()
     database.database_name = "test_database"
-    database.get_all_table_names_in_schema.return_value = [
-        DatasourceName("table1", "schema1"),
-        DatasourceName("table2", "schema1"),
-    ]
-    database.get_all_view_names_in_schema.return_value = [
-        DatasourceName("view1", "schema1"),
-    ]
+    database.get_all_table_names_in_schema.return_value = {
+        ("table1", "schema1", None),
+        ("table2", "schema1", None),
+    }
+    database.get_all_view_names_in_schema.return_value = {
+        ("view1", "schema1", None),
+    }
 
     DatabaseDAO = mocker.patch("superset.commands.database.tables.DatabaseDAO")  # noqa: N806
     DatabaseDAO.find_by_id.return_value = database
diff --git a/tests/unit_tests/models/core_test.py b/tests/unit_tests/models/core_test.py
index f399fc852..5b269fc3b 100644
--- a/tests/unit_tests/models/core_test.py
+++ b/tests/unit_tests/models/core_test.py
@@ -768,3 +768,59 @@ FROM (
 WHERE
   TRUE AND TRUE"""
     )
+
+
+def test_get_all_table_names_in_schema(mocker: MockerFixture) -> None:
+    """
+    Test the `get_all_table_names_in_schema` method.
+    """
+    database = Database(
+        database_name="db",
+        sqlalchemy_uri="postgresql://user:password@host:5432/examples",
+    )
+
+    mocker.patch.object(database, "get_inspector")
+    get_table_names = mocker.patch(
+        "superset.db_engine_specs.postgres.PostgresEngineSpec.get_table_names"
+    )
+    get_table_names.return_value = {"first_table", "second_table", "third_table"}
+
+    tables_list = database.get_all_table_names_in_schema(
+        catalog="examples",
+        schema="public",
+    )
+    assert sorted(tables_list) == sorted(
+        {
+            ("first_table", "public", "examples"),
+            ("second_table", "public", "examples"),
+            ("third_table", "public", "examples"),
+        }
+    )
+
+
+def test_get_all_view_names_in_schema(mocker: MockerFixture) -> None:
+    """
+    Test the `get_all_view_names_in_schema` method.
+    """
+    database = Database(
+        database_name="db",
+        sqlalchemy_uri="postgresql://user:password@host:5432/examples",
+    )
+
+    mocker.patch.object(database, "get_inspector")
+    get_view_names = mocker.patch(
+        "superset.db_engine_specs.base.BaseEngineSpec.get_view_names"
+    )
+    get_view_names.return_value = {"first_view", "second_view", "third_view"}
+
+    views_list = database.get_all_view_names_in_schema(
+        catalog="examples",
+        schema="public",
+    )
+    assert sorted(views_list) == sorted(
+        {
+            ("first_view", "public", "examples"),
+            ("second_view", "public", "examples"),
+            ("third_view", "public", "examples"),
+        }
+    )