diff --git a/superset/views/base_api.py b/superset/views/base_api.py index 5a59b8014..d507d4ace 100644 --- a/superset/views/base_api.py +++ b/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/", methods=["GET"]) @protect() @safe diff --git a/tests/css_templates/api_tests.py b/tests/css_templates/api_tests.py index 2d4073f63..23a6c7559 100644 --- a/tests/css_templates/api_tests.py +++ b/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") diff --git a/tests/databases/api_tests.py b/tests/databases/api_tests.py index fbf1b9162..62a140059 100644 --- a/tests/databases/api_tests.py +++ b/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 diff --git a/tests/datasets/api_tests.py b/tests/datasets/api_tests.py index e2ad09a29..43874299f 100644 --- a/tests/datasets/api_tests.py +++ b/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