feat: add statsd metrics to FAB rest API post, put and delete (#11362)

superset/views/base_api.py

@@ -319,6 +319,30 @@ class BaseSupersetModelRestApi(ModelRestApi):
         self.send_stats_metrics(response, self.get_list.__name__, duration)
         return response
 
+    def post_headless(self) -> Response:
+        """
+        Add statsd metrics to builtin FAB POST endpoint
+        """
+        duration, response = time_function(super().post_headless)
+        self.send_stats_metrics(response, self.post.__name__, duration)
+        return response
+
+    def put_headless(self, pk: int) -> Response:
+        """
+        Add statsd metrics to builtin FAB PUT endpoint
+        """
+        duration, response = time_function(super().put_headless, pk)
+        self.send_stats_metrics(response, self.put.__name__, duration)
+        return response
+
+    def delete_headless(self, pk: int) -> Response:
+        """
+        Add statsd metrics to builtin FAB DELETE endpoint
+        """
+        duration, response = time_function(super().delete_headless, pk)
+        self.send_stats_metrics(response, self.delete.__name__, duration)
+        return response
+
     @expose("/related/<column_name>", methods=["GET"])
     @protect()
     @safe

tests/css_templates/api_tests.py

@@ -196,7 +196,7 @@ class TestCssTemplateApi(SupersetTestCase):
         max_id = db.session.query(func.max(CssTemplate.id)).scalar()
         self.login(username="admin")
         uri = f"api/v1/css_template/{max_id + 1}"
-        rv = self.client.get(uri)
+        rv = self.get_assert_metric(uri, "get")
         assert rv.status_code == 404
 
     def test_create_css_template(self):
@@ -210,7 +210,7 @@ class TestCssTemplateApi(SupersetTestCase):
 
         self.login(username="admin")
         uri = f"api/v1/css_template/"
-        rv = self.client.post(uri, json=post_data)
+        rv = self.post_assert_metric(uri, post_data, "post")
         data = json.loads(rv.data.decode("utf-8"))
         assert rv.status_code == 201
 
@@ -241,7 +241,7 @@ class TestCssTemplateApi(SupersetTestCase):
 
         self.login(username="admin")
         uri = f"api/v1/css_template/{css_template.id}"
-        rv = self.client.put(uri, json=put_data)
+        rv = self.put_assert_metric(uri, put_data, "put")
         assert rv.status_code == 200
 
         model = db.session.query(CssTemplate).get(css_template.id)
@@ -262,7 +262,7 @@ class TestCssTemplateApi(SupersetTestCase):
         }
 
         uri = f"api/v1/css_template/{max_id + 1}"
-        rv = self.client.put(uri, json=put_data)
+        rv = self.put_assert_metric(uri, put_data, "put")
         assert rv.status_code == 404
 
     @pytest.mark.usefixtures("create_css_templates")
@@ -278,7 +278,7 @@ class TestCssTemplateApi(SupersetTestCase):
 
         self.login(username="admin")
         uri = f"api/v1/css_template/{css_template.id}"
-        rv = self.client.delete(uri)
+        rv = self.delete_assert_metric(uri, "delete")
         assert rv.status_code == 200
 
         model = db.session.query(CssTemplate).get(css_template.id)
@@ -292,7 +292,7 @@ class TestCssTemplateApi(SupersetTestCase):
         max_id = db.session.query(func.max(CssTemplate.id)).scalar()
         self.login(username="admin")
         uri = f"api/v1/css_template/{max_id + 1}"
-        rv = self.client.delete(uri)
+        rv = self.delete_assert_metric(uri, "delete")
         assert rv.status_code == 404
 
     @pytest.mark.usefixtures("create_css_templates")

tests/databases/api_tests.py

@@ -795,13 +795,13 @@ class TestDatabaseApi(SupersetTestCase):
         invalid_id = max_id + 1
         uri = f"api/v1/database/{invalid_id}/related_objects/"
         self.login(username="admin")
-        rv = self.client.get(uri)
+        rv = self.get_assert_metric(uri, "related_objects")
         self.assertEqual(rv.status_code, 404)
         self.logout()
         self.login(username="gamma")
         database = get_example_database()
         uri = f"api/v1/database/{database.id}/related_objects/"
-        rv = self.client.get(uri)
+        rv = self.get_assert_metric(uri, "related_objects")
         self.assertEqual(rv.status_code, 404)
 
     def test_export_database(self):
@@ -812,8 +812,7 @@ class TestDatabaseApi(SupersetTestCase):
         database = get_example_database()
         argument = [database.id]
         uri = f"api/v1/database/export/?q={prison.dumps(argument)}"
-        rv = self.client.get(uri)
-
+        rv = self.get_assert_metric(uri, "export")
         assert rv.status_code == 200
 
         buf = BytesIO(rv.data)
@@ -828,7 +827,6 @@ class TestDatabaseApi(SupersetTestCase):
         argument = [database.id]
         uri = f"api/v1/database/export/?q={prison.dumps(argument)}"
         rv = self.client.get(uri)
-
         assert rv.status_code == 401
 
     def test_export_database_non_existing(self):
@@ -842,6 +840,5 @@ class TestDatabaseApi(SupersetTestCase):
         self.login(username="admin")
         argument = [invalid_id]
         uri = f"api/v1/database/export/?q={prison.dumps(argument)}"
-        rv = self.client.get(uri)
-
+        rv = self.get_assert_metric(uri, "export")
         assert rv.status_code == 404

tests/datasets/api_tests.py

@@ -110,9 +110,9 @@ class TestDatasetApi(SupersetTestCase):
         }
         uri = f"api/v1/dataset/?q={prison.dumps(arguments)}"
         rv = self.get_assert_metric(uri, "get_list")
-        self.assertEqual(rv.status_code, 200)
+        assert rv.status_code == 200
         response = json.loads(rv.data.decode("utf-8"))
-        self.assertEqual(response["count"], 1)
+        assert response["count"] == 1
         expected_columns = [
             "changed_by",
             "changed_by_name",
@@ -129,7 +129,7 @@ class TestDatasetApi(SupersetTestCase):
             "sql",
             "table_name",
         ]
-        self.assertEqual(sorted(list(response["result"][0].keys())), expected_columns)
+        assert sorted(list(response["result"][0].keys())) == expected_columns
 
     def test_get_dataset_list_gamma(self):
         """
@@ -139,9 +139,9 @@ class TestDatasetApi(SupersetTestCase):
         self.login(username="gamma")
         uri = "api/v1/dataset/"
         rv = self.get_assert_metric(uri, "get_list")
-        self.assertEqual(rv.status_code, 200)
+        assert rv.status_code == 200
         response = json.loads(rv.data.decode("utf-8"))
-        self.assertEqual(response["result"], [])
+        assert response["result"] == []
 
     def test_get_dataset_related_database_gamma(self):
         """
@@ -150,10 +150,10 @@ class TestDatasetApi(SupersetTestCase):
         self.login(username="gamma")
         uri = "api/v1/dataset/related/database"
         rv = self.client.get(uri)
-        self.assertEqual(rv.status_code, 200)
+        assert rv.status_code == 200
         response = json.loads(rv.data.decode("utf-8"))
-        self.assertEqual(response["count"], 0)
-        self.assertEqual(response["result"], [])
+        assert response["count"] == 0
+        assert response["result"] == []
 
     def test_get_dataset_item(self):
         """
@@ -163,7 +163,7 @@ class TestDatasetApi(SupersetTestCase):
         self.login(username="admin")
         uri = f"api/v1/dataset/{table.id}"
         rv = self.get_assert_metric(uri, "get")
-        self.assertEqual(rv.status_code, 200)
+        assert rv.status_code == 200
         response = json.loads(rv.data.decode("utf-8"))
         expected_result = {
             "cache_timeout": None,
@@ -184,8 +184,8 @@ class TestDatasetApi(SupersetTestCase):
         assert {
             k: v for k, v in response["result"].items() if k in expected_result
         } == expected_result
-        self.assertEqual(len(response["result"]["columns"]), 3)
-        self.assertEqual(len(response["result"]["metrics"]), 2)
+        assert len(response["result"]["columns"]) == 3
+        assert len(response["result"]["metrics"]) == 2
 
     def test_get_dataset_distinct_schema(self):
         """
@@ -196,8 +196,8 @@ class TestDatasetApi(SupersetTestCase):
             uri = f"api/v1/dataset/distinct/schema?q={prison.dumps(query_parameter)}"
             rv = self.client.get(uri)
             response = json.loads(rv.data.decode("utf-8"))
-            self.assertEqual(rv.status_code, 200)
-            self.assertEqual(response, expected_response)
+            assert rv.status_code == 200
+            assert response == expected_response
 
         example_db = get_example_database()
         datasets = []
@@ -224,8 +224,8 @@ class TestDatasetApi(SupersetTestCase):
             uri = "api/v1/dataset/distinct/schema"
             rv = self.client.get(uri)
             response = json.loads(rv.data.decode("utf-8"))
-            self.assertEqual(rv.status_code, 200)
-            self.assertEqual(response, expected_response)
+            assert rv.status_code == 200
+            assert response == expected_response
 
             # Test filter
             query_parameter = {"filter": "inf"}
@@ -264,7 +264,7 @@ class TestDatasetApi(SupersetTestCase):
         self.login(username="admin")
         uri = "api/v1/dataset/distinct/table_name"
         rv = self.client.get(uri)
-        self.assertEqual(rv.status_code, 404)
+        assert rv.status_code == 404
 
     def test_get_dataset_distinct_gamma(self):
         """
@@ -275,10 +275,10 @@ class TestDatasetApi(SupersetTestCase):
         self.login(username="gamma")
         uri = "api/v1/dataset/distinct/schema"
         rv = self.client.get(uri)
-        self.assertEqual(rv.status_code, 200)
+        assert rv.status_code == 200
         response = json.loads(rv.data.decode("utf-8"))
-        self.assertEqual(response["count"], 0)
-        self.assertEqual(response["result"], [])
+        assert response["count"] == 0
+        assert response["result"] == []
 
         db.session.delete(dataset)
         db.session.commit()
@@ -290,7 +290,7 @@ class TestDatasetApi(SupersetTestCase):
         self.login(username="admin")
         uri = "api/v1/dataset/_info"
         rv = self.get_assert_metric(uri, "info")
-        self.assertEqual(rv.status_code, 200)
+        assert rv.status_code == 200
 
     def test_create_dataset_item(self):
         """
@@ -305,12 +305,12 @@ class TestDatasetApi(SupersetTestCase):
         }
         uri = "api/v1/dataset/"
         rv = self.post_assert_metric(uri, table_data, "post")
-        self.assertEqual(rv.status_code, 201)
+        assert rv.status_code == 201
         data = json.loads(rv.data.decode("utf-8"))
         table_id = data.get("id")
         model = db.session.query(SqlaTable).get(table_id)
-        self.assertEqual(model.table_name, table_data["table_name"])
-        self.assertEqual(model.database_id, table_data["database"])
+        assert model.table_name == table_data["table_name"]
+        assert model.database_id == table_data["database"]
 
         # Assert that columns were created
         columns = (
@@ -319,8 +319,8 @@ class TestDatasetApi(SupersetTestCase):
             .order_by("column_name")
             .all()
         )
-        self.assertEqual(columns[0].column_name, "id")
-        self.assertEqual(columns[1].column_name, "name")
+        assert columns[0].column_name == "id"
+        assert columns[1].column_name == "name"
 
         # Assert that metrics were created
         columns = (
@@ -329,7 +329,7 @@ class TestDatasetApi(SupersetTestCase):
             .order_by("metric_name")
             .all()
         )
-        self.assertEqual(columns[0].expression, "COUNT(*)")
+        assert columns[0].expression == "COUNT(*)"
 
         db.session.delete(model)
         db.session.commit()
@@ -347,7 +347,7 @@ class TestDatasetApi(SupersetTestCase):
         }
         uri = "api/v1/dataset/"
         rv = self.client.post(uri, json=table_data)
-        self.assertEqual(rv.status_code, 401)
+        assert rv.status_code == 401
 
     def test_create_dataset_item_owner(self):
         """
@@ -366,11 +366,11 @@ class TestDatasetApi(SupersetTestCase):
         }
         uri = "api/v1/dataset/"
         rv = self.post_assert_metric(uri, table_data, "post")
-        self.assertEqual(rv.status_code, 201)
+        assert rv.status_code == 201
         data = json.loads(rv.data.decode("utf-8"))
         model = db.session.query(SqlaTable).get(data.get("id"))
-        self.assertIn(admin, model.owners)
-        self.assertIn(alpha, model.owners)
+        assert admin in model.owners
+        assert alpha in model.owners
         db.session.delete(model)
         db.session.commit()
 
@@ -389,10 +389,10 @@ class TestDatasetApi(SupersetTestCase):
         }
         uri = f"api/v1/dataset/"
         rv = self.post_assert_metric(uri, table_data, "post")
-        self.assertEqual(rv.status_code, 422)
+        assert rv.status_code == 422
         data = json.loads(rv.data.decode("utf-8"))
         expected_result = {"message": {"owners": ["Owners are invalid"]}}
-        self.assertEqual(data, expected_result)
+        assert data == expected_result
 
     def test_create_dataset_validate_uniqueness(self):
         """
@@ -407,11 +407,11 @@ class TestDatasetApi(SupersetTestCase):
         }
         uri = "api/v1/dataset/"
         rv = self.post_assert_metric(uri, table_data, "post")
-        self.assertEqual(rv.status_code, 422)
+        assert rv.status_code == 422
         data = json.loads(rv.data.decode("utf-8"))
-        self.assertEqual(
-            data, {"message": {"table_name": ["Datasource birth_names already exists"]}}
-        )
+        assert data == {
+            "message": {"table_name": ["Datasource birth_names already exists"]}
+        }
 
     def test_create_dataset_same_name_different_schema(self):
         if backend() == "sqlite":
@@ -432,13 +432,13 @@ class TestDatasetApi(SupersetTestCase):
 
         uri = "api/v1/dataset/"
         rv = self.post_assert_metric(uri, table_data, "post")
-        self.assertEqual(rv.status_code, 201)
+        assert rv.status_code == 201
 
         # cleanup
         data = json.loads(rv.data.decode("utf-8"))
         uri = f'api/v1/dataset/{data.get("id")}'
         rv = self.client.delete(uri)
-        self.assertEqual(rv.status_code, 200)
+        assert rv.status_code == 200
         example_db.get_sqla_engine().execute(
             f"DROP TABLE {CTAS_SCHEMA_NAME}.birth_names"
         )
@@ -451,9 +451,9 @@ class TestDatasetApi(SupersetTestCase):
         dataset_data = {"database": 1000, "schema": "", "table_name": "birth_names"}
         uri = "api/v1/dataset/"
         rv = self.post_assert_metric(uri, dataset_data, "post")
-        self.assertEqual(rv.status_code, 422)
+        assert rv.status_code == 422
         data = json.loads(rv.data.decode("utf-8"))
-        self.assertEqual(data, {"message": {"database": ["Database does not exist"]}})
+        assert data == {"message": {"database": ["Database does not exist"]}}
 
     def test_create_dataset_validate_tables_exists(self):
         """
@@ -468,7 +468,7 @@ class TestDatasetApi(SupersetTestCase):
         }
         uri = "api/v1/dataset/"
         rv = self.post_assert_metric(uri, table_data, "post")
-        self.assertEqual(rv.status_code, 422)
+        assert rv.status_code == 422
 
     @patch("superset.datasets.dao.DatasetDAO.create")
     def test_create_dataset_sqlalchemy_error(self, mock_dao_create):
@@ -486,8 +486,8 @@ class TestDatasetApi(SupersetTestCase):
         uri = "api/v1/dataset/"
         rv = self.post_assert_metric(uri, dataset_data, "post")
         data = json.loads(rv.data.decode("utf-8"))
-        self.assertEqual(rv.status_code, 422)
-        self.assertEqual(data, {"message": "Dataset could not be created."})
+        assert rv.status_code == 422
+        assert data == {"message": "Dataset could not be created."}
 
     def test_update_dataset_item(self):
         """
@@ -498,9 +498,9 @@ class TestDatasetApi(SupersetTestCase):
         dataset_data = {"description": "changed_description"}
         uri = f"api/v1/dataset/{dataset.id}"
         rv = self.put_assert_metric(uri, dataset_data, "put")
-        self.assertEqual(rv.status_code, 200)
+        assert rv.status_code == 200
         model = db.session.query(SqlaTable).get(dataset.id)
-        self.assertEqual(model.description, dataset_data["description"])
+        assert model.description == dataset_data["description"]
 
         db.session.delete(dataset)
         db.session.commit()
@@ -532,7 +532,7 @@ class TestDatasetApi(SupersetTestCase):
         data["result"]["columns"].append(new_column_data)
         rv = self.client.put(uri, json={"columns": data["result"]["columns"]})
 
-        self.assertEqual(rv.status_code, 200)
+        assert rv.status_code == 200
 
         columns = (
             db.session.query(TableColumn)
@@ -540,13 +540,13 @@ class TestDatasetApi(SupersetTestCase):
             .order_by("column_name")
             .all()
         )
-        self.assertEqual(columns[0].column_name, "id")
-        self.assertEqual(columns[1].column_name, "name")
-        self.assertEqual(columns[2].column_name, new_column_data["column_name"])
-        self.assertEqual(columns[2].description, new_column_data["description"])
-        self.assertEqual(columns[2].expression, new_column_data["expression"])
-        self.assertEqual(columns[2].type, new_column_data["type"])
-        self.assertEqual(columns[2].verbose_name, new_column_data["verbose_name"])
+        assert columns[0].column_name == "id"
+        assert columns[1].column_name == "name"
+        assert columns[2].column_name == new_column_data["column_name"]
+        assert columns[2].description == new_column_data["description"]
+        assert columns[2].expression == new_column_data["expression"]
+        assert columns[2].type == new_column_data["type"]
+        assert columns[2].verbose_name == new_column_data["verbose_name"]
 
         db.session.delete(dataset)
         db.session.commit()
@@ -569,19 +569,19 @@ class TestDatasetApi(SupersetTestCase):
         resp_columns[0]["groupby"] = False
         resp_columns[0]["filterable"] = False
         rv = self.client.put(uri, json={"columns": resp_columns})
-        self.assertEqual(rv.status_code, 200)
+        assert rv.status_code == 200
         columns = (
             db.session.query(TableColumn)
             .filter_by(table_id=dataset.id)
             .order_by("column_name")
             .all()
         )
-        self.assertEqual(columns[0].column_name, "id")
-        self.assertEqual(columns[1].column_name, "name")
+        assert columns[0].column_name == "id"
+        assert columns[1].column_name, "name"
         # TODO(bkyryliuk): find the reason why update is failing for the presto database
         if get_example_database().backend != "presto":
-            self.assertEqual(columns[0].groupby, False)
-            self.assertEqual(columns[0].filterable, False)
+            assert columns[0].groupby is False
+            assert columns[0].filterable is False
 
         db.session.delete(dataset)
         db.session.commit()
@@ -597,12 +597,12 @@ class TestDatasetApi(SupersetTestCase):
         # try to insert a new column ID that already exists
         data = {"columns": [{"column_name": "id", "type": "INTEGER"}]}
         rv = self.put_assert_metric(uri, data, "put")
-        self.assertEqual(rv.status_code, 422)
+        assert rv.status_code == 422
         data = json.loads(rv.data.decode("utf-8"))
         expected_result = {
             "message": {"columns": ["One or more columns already exist"]}
         }
-        self.assertEqual(data, expected_result)
+        assert data == expected_result
         db.session.delete(dataset)
         db.session.commit()
 
@@ -617,12 +617,12 @@ class TestDatasetApi(SupersetTestCase):
         # try to insert a new column ID that already exists
         data = {"metrics": [{"metric_name": "count", "expression": "COUNT(*)"}]}
         rv = self.put_assert_metric(uri, data, "put")
-        self.assertEqual(rv.status_code, 422)
+        assert rv.status_code == 422
         data = json.loads(rv.data.decode("utf-8"))
         expected_result = {
             "message": {"metrics": ["One or more metrics already exist"]}
         }
-        self.assertEqual(data, expected_result)
+        assert data == expected_result
         db.session.delete(dataset)
         db.session.commit()
 
@@ -642,12 +642,12 @@ class TestDatasetApi(SupersetTestCase):
             ]
         }
         rv = self.put_assert_metric(uri, data, "put")
-        self.assertEqual(rv.status_code, 422)
+        assert rv.status_code == 422
         data = json.loads(rv.data.decode("utf-8"))
         expected_result = {
             "message": {"columns": ["One or more columns are duplicated"]}
         }
-        self.assertEqual(data, expected_result)
+        assert data == expected_result
         db.session.delete(dataset)
         db.session.commit()
 
@@ -667,12 +667,12 @@ class TestDatasetApi(SupersetTestCase):
             ]
         }
         rv = self.put_assert_metric(uri, data, "put")
-        self.assertEqual(rv.status_code, 422)
+        assert rv.status_code == 422
         data = json.loads(rv.data.decode("utf-8"))
         expected_result = {
             "message": {"metrics": ["One or more metrics are duplicated"]}
         }
-        self.assertEqual(data, expected_result)
+        assert data == expected_result
         db.session.delete(dataset)
         db.session.commit()
 
@@ -685,7 +685,7 @@ class TestDatasetApi(SupersetTestCase):
         table_data = {"description": "changed_description"}
         uri = f"api/v1/dataset/{dataset.id}"
         rv = self.client.put(uri, json=table_data)
-        self.assertEqual(rv.status_code, 401)
+        assert rv.status_code == 401
         db.session.delete(dataset)
         db.session.commit()
 
@@ -698,7 +698,7 @@ class TestDatasetApi(SupersetTestCase):
         table_data = {"description": "changed_description"}
         uri = f"api/v1/dataset/{dataset.id}"
         rv = self.put_assert_metric(uri, table_data, "put")
-        self.assertEqual(rv.status_code, 403)
+        assert rv.status_code == 403
         db.session.delete(dataset)
         db.session.commit()
 
@@ -711,7 +711,7 @@ class TestDatasetApi(SupersetTestCase):
         table_data = {"description": "changed_description", "owners": [1000]}
         uri = f"api/v1/dataset/{dataset.id}"
         rv = self.put_assert_metric(uri, table_data, "put")
-        self.assertEqual(rv.status_code, 422)
+        assert rv.status_code == 422
         db.session.delete(dataset)
         db.session.commit()
 
@@ -728,11 +728,11 @@ class TestDatasetApi(SupersetTestCase):
         uri = f"api/v1/dataset/{dataset.id}"
         rv = self.put_assert_metric(uri, table_data, "put")
         data = json.loads(rv.data.decode("utf-8"))
-        self.assertEqual(rv.status_code, 422)
+        assert rv.status_code == 422
         expected_response = {
             "message": {"table_name": ["Datasource ab_user already exists"]}
         }
-        self.assertEqual(data, expected_response)
+        assert data == expected_response
         db.session.delete(dataset)
         db.session.delete(ab_user)
         db.session.commit()
@@ -750,8 +750,8 @@ class TestDatasetApi(SupersetTestCase):
         uri = f"api/v1/dataset/{dataset.id}"
         rv = self.client.put(uri, json=table_data)
         data = json.loads(rv.data.decode("utf-8"))
-        self.assertEqual(rv.status_code, 422)
-        self.assertEqual(data, {"message": "Dataset could not be updated."})
+        assert rv.status_code == 422
+        assert data == {"message": "Dataset could not be updated."}
 
         db.session.delete(dataset)
         db.session.commit()
@@ -762,16 +762,16 @@ class TestDatasetApi(SupersetTestCase):
         """
         dataset = self.insert_default_dataset()
         view_menu = security_manager.find_view_menu(dataset.get_perm())
-        self.assertIsNotNone(view_menu)
+        assert view_menu is not None
         view_menu_id = view_menu.id
         self.login(username="admin")
         uri = f"api/v1/dataset/{dataset.id}"
         rv = self.client.delete(uri)
-        self.assertEqual(rv.status_code, 200)
+        assert rv.status_code == 200
         non_view_menu = db.session.query(security_manager.viewmenu_model).get(
             view_menu_id
         )
-        self.assertIsNone(non_view_menu)
+        assert non_view_menu is None
 
     def test_delete_item_dataset_not_owned(self):
         """
@@ -781,7 +781,7 @@ class TestDatasetApi(SupersetTestCase):
         self.login(username="alpha")
         uri = f"api/v1/dataset/{dataset.id}"
         rv = self.delete_assert_metric(uri, "delete")
-        self.assertEqual(rv.status_code, 403)
+        assert rv.status_code == 403
         db.session.delete(dataset)
         db.session.commit()
 
@@ -793,7 +793,7 @@ class TestDatasetApi(SupersetTestCase):
         self.login(username="gamma")
         uri = f"api/v1/dataset/{dataset.id}"
         rv = self.client.delete(uri)
-        self.assertEqual(rv.status_code, 401)
+        assert rv.status_code == 401
         db.session.delete(dataset)
         db.session.commit()
 
@@ -809,8 +809,8 @@ class TestDatasetApi(SupersetTestCase):
         uri = f"api/v1/dataset/{dataset.id}"
         rv = self.delete_assert_metric(uri, "delete")
         data = json.loads(rv.data.decode("utf-8"))
-        self.assertEqual(rv.status_code, 422)
-        self.assertEqual(data, {"message": "Dataset could not be deleted."})
+        assert rv.status_code == 422
+        assert data == {"message": "Dataset could not be deleted."}
         db.session.delete(dataset)
         db.session.commit()
 
@@ -914,14 +914,14 @@ class TestDatasetApi(SupersetTestCase):
         self.login(username="admin")
         uri = f"api/v1/dataset/{dataset.id}/refresh"
         rv = self.put_assert_metric(uri, {}, "refresh")
-        self.assertEqual(rv.status_code, 200)
+        assert rv.status_code == 200
         # Assert the column is restored on refresh
         id_column = (
             db.session.query(TableColumn)
             .filter_by(table_id=dataset.id, column_name="id")
             .one()
         )
-        self.assertIsNotNone(id_column)
+        assert id_column is not None
         db.session.delete(dataset)
         db.session.commit()
 
@@ -934,7 +934,7 @@ class TestDatasetApi(SupersetTestCase):
         self.login(username="admin")
         uri = f"api/v1/dataset/{max_id + 1}/refresh"
         rv = self.put_assert_metric(uri, {}, "refresh")
-        self.assertEqual(rv.status_code, 404)
+        assert rv.status_code == 404
 
     def test_dataset_item_refresh_not_owned(self):
         """
@@ -944,7 +944,7 @@ class TestDatasetApi(SupersetTestCase):
         self.login(username="alpha")
         uri = f"api/v1/dataset/{dataset.id}/refresh"
         rv = self.put_assert_metric(uri, {}, "refresh")
-        self.assertEqual(rv.status_code, 403)
+        assert rv.status_code == 403
 
         db.session.delete(dataset)
         db.session.commit()
@@ -964,7 +964,7 @@ class TestDatasetApi(SupersetTestCase):
 
         self.login(username="admin")
         rv = self.get_assert_metric(uri, "export")
-        self.assertEqual(rv.status_code, 200)
+        assert rv.status_code == 200
 
         cli_export = export_to_dict(
             session=db.session,
@@ -979,7 +979,7 @@ class TestDatasetApi(SupersetTestCase):
                 expected_response = export_table
                 break
         ui_export = yaml.safe_load(rv.data.decode("utf-8"))
-        self.assertEqual(ui_export[0], expected_response)
+        assert ui_export[0] == expected_response
 
     def test_export_dataset_not_found(self):
         """
@@ -991,7 +991,7 @@ class TestDatasetApi(SupersetTestCase):
         uri = f"api/v1/dataset/export/?q={prison.dumps(argument)}"
         self.login(username="admin")
         rv = self.get_assert_metric(uri, "export")
-        self.assertEqual(rv.status_code, 404)
+        assert rv.status_code == 404
 
     def test_export_dataset_gamma(self):
         """
@@ -1004,7 +1004,7 @@ class TestDatasetApi(SupersetTestCase):
 
         self.login(username="gamma")
         rv = self.client.get(uri)
-        self.assertEqual(rv.status_code, 401)
+        assert rv.status_code == 401
 
     def test_get_dataset_related_objects(self):
         """
@@ -1015,10 +1015,10 @@ class TestDatasetApi(SupersetTestCase):
         table = self.get_birth_names_dataset()
         uri = f"api/v1/dataset/{table.id}/related_objects"
         rv = self.get_assert_metric(uri, "related_objects")
-        self.assertEqual(rv.status_code, 200)
         response = json.loads(rv.data.decode("utf-8"))
-        self.assertEqual(response["charts"]["count"], 18)
-        self.assertEqual(response["dashboards"]["count"], 1)
+        assert rv.status_code == 200
+        assert response["charts"]["count"] == 18
+        assert response["dashboards"]["count"] == 1
 
     def test_get_dataset_related_objects_not_found(self):
         """
@@ -1030,10 +1030,10 @@ class TestDatasetApi(SupersetTestCase):
         uri = f"api/v1/dataset/{invalid_id}/related_objects/"
         self.login(username="admin")
         rv = self.client.get(uri)
-        self.assertEqual(rv.status_code, 404)
+        assert rv.status_code == 404
         self.logout()
         self.login(username="gamma")
         table = self.get_birth_names_dataset()
         uri = f"api/v1/dataset/{table.id}/related_objects"
         rv = self.client.get(uri)
-        self.assertEqual(rv.status_code, 404)
+        assert rv.status_code == 404