fix test

apache · Jun 16, 2022 · ed31bf1 · ed31bf1
1 parent 0653a58
commit ed31bf1
Showing 1 changed file with 15 additions and 6 deletions.
diff --git a/tests/integration_tests/datasource_tests.py b/tests/integration_tests/datasource_tests.py
@@ -392,16 +392,25 @@ def my_check(datasource):
         app.config["DATASET_HEALTH_CHECK"] = None
 
     def test_get_datasource_failed(self):
+        from superset.datasource.dao import DatasourceDAO
+
         pytest.raises(
             DatasourceNotFound,
-            lambda: DatasourceDAO.get_datasource("table", 9999999, db.session),
+            lambda: DatasourceDAO.get_datasource(db.session, "table", 9999999),
         )
 
         self.login(username="admin")
-        resp = self.get_json_resp("/datasource/get/druid/500000/", raise_on_error=False)
-        self.assertEqual(resp.get("error"), "Dataset does not exist")
+        resp = self.get_json_resp("/datasource/get/table/500000/", raise_on_error=False)
+        self.assertEqual(resp.get("error"), "Datasource does not exist")
 
-        resp = self.get_json_resp(
-            "/datasource/get/invalid-datasource-type/500000/", raise_on_error=False
+    def test_get_datasource_invalid_datasource_failed(self):
+        from superset.datasource.dao import DatasourceDAO
+
+        pytest.raises(
+            DatasourceTypeNotSupportedError,
+            lambda: DatasourceDAO.get_datasource(db.session, "druid", 9999999),
         )
-        self.assertEqual(resp.get("error"), "Dataset does not exist")
+
+        self.login(username="admin")
+        resp = self.get_json_resp("/datasource/get/druid/500000/", raise_on_error=False)
+        self.assertEqual(resp.get("error"), "'druid' is not a valid DatasourceType")