From 5146afa83217c5468efc6b525cd702e94dc353ff Mon Sep 17 00:00:00 2001 From: Grace Date: Tue, 15 Dec 2020 15:30:45 -0800 Subject: [PATCH] add unit test --- superset/connectors/sqla/models.py | 2 +- tests/datasource_tests.py | 18 +++++++++++++++++- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py index 6214a984e700c..d8b912b239cc9 100644 --- a/superset/connectors/sqla/models.py +++ b/superset/connectors/sqla/models.py @@ -1474,7 +1474,7 @@ class and any keys added via `ExtraCache`. return extra_cache_keys def health_check(self, commit: bool = False, force: bool = False) -> None: - check = config["DATASET_HEALTH_CHECK"] + check = config.get("DATASET_HEALTH_CHECK") if check is None: return diff --git a/tests/datasource_tests.py b/tests/datasource_tests.py index 31a4c96334188..890b4a63e42c6 100644 --- a/tests/datasource_tests.py +++ b/tests/datasource_tests.py @@ -18,7 +18,7 @@ import json from copy import deepcopy -from superset import db +from superset import app, db from superset.connectors.sqla.models import SqlaTable from superset.utils.core import get_example_database @@ -190,6 +190,22 @@ def test_get_datasource(self): }, ) + def test_get_datasource_with_health_check(self): + def my_check(datasource): + return "Warning message!" + + app.config["DATASET_HEALTH_CHECK"] = my_check + my_check.version = 0.1 + + self.login(username="admin") + tbl = self.get_table_by_name("birth_names") + url = f"/datasource/get/{tbl.type}/{tbl.id}/" + tbl.health_check(commit=True, force=True) + resp = self.get_json_resp(url) + self.assertEqual(resp["health_check_message"], "Warning message!") + + del app.config["DATASET_HEALTH_CHECK"] + def test_get_datasource_failed(self): self.login(username="admin") url = f"/datasource/get/druid/500000/"