From 595e41da9953a9978bedd133af1a5b635a67be75 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Fri, 5 Aug 2022 11:42:55 -0700
Subject: [PATCH] Add unit tests

---
 tests/unit_tests/models/__init__.py  | 16 +++++++
 tests/unit_tests/models/core_test.py | 72 ++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+)
 create mode 100644 tests/unit_tests/models/__init__.py
 create mode 100644 tests/unit_tests/models/core_test.py

diff --git a/tests/unit_tests/models/__init__.py b/tests/unit_tests/models/__init__.py
new file mode 100644
index 0000000000000..13a83393a9124
--- /dev/null
+++ b/tests/unit_tests/models/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/tests/unit_tests/models/core_test.py b/tests/unit_tests/models/core_test.py
new file mode 100644
index 0000000000000..3338ddcb61441
--- /dev/null
+++ b/tests/unit_tests/models/core_test.py
@@ -0,0 +1,72 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# pylint: disable=import-outside-toplevel
+
+from typing import List, Optional
+
+from pytest_mock import MockFixture
+from sqlalchemy.engine.reflection import Inspector
+
+
+def test_get_metrics(mocker: MockFixture) -> None:
+    """
+    Tests for ``get_metrics``.
+    """
+    from superset.db_engine_specs.base import MetricType
+    from superset.db_engine_specs.sqlite import SqliteEngineSpec
+    from superset.models.core import Database
+
+    database = Database(database_name="my_database", sqlalchemy_uri="sqlite://")
+    assert database.get_metrics("table") == [
+        {
+            "expression": "COUNT(*)",
+            "metric_name": "count",
+            "metric_type": "count",
+            "verbose_name": "COUNT(*)",
+        }
+    ]
+
+    class CustomSqliteEngineSpec(SqliteEngineSpec):
+        @classmethod
+        def get_metrics(
+            cls,
+            database: Database,
+            inspector: Inspector,
+            table_name: str,
+            schema: Optional[str],
+        ) -> List[MetricType]:
+            return [
+                {
+                    "expression": "COUNT(DISTINCT user_id)",
+                    "metric_name": "count_distinct_user_id",
+                    "metric_type": "count_distinct",
+                    "verbose_name": "COUNT(DISTINCT user_id)",
+                },
+            ]
+
+    database.get_db_engine_spec_for_backend = mocker.MagicMock(  # type: ignore
+        return_value=CustomSqliteEngineSpec
+    )
+    assert database.get_metrics("table") == [
+        {
+            "expression": "COUNT(DISTINCT user_id)",
+            "metric_name": "count_distinct_user_id",
+            "metric_type": "count_distinct",
+            "verbose_name": "COUNT(DISTINCT user_id)",
+        },
+    ]