diff --git a/tests/integration_tests/datasets/api_tests.py b/tests/integration_tests/datasets/api_tests.py index 9208e7a8ee4d2..c5624230335a3 100644 --- a/tests/integration_tests/datasets/api_tests.py +++ b/tests/integration_tests/datasets/api_tests.py @@ -786,20 +786,39 @@ def test_update_dataset_item_w_override_columns_same_columns(self): # Add default dataset main_db = get_main_database() - dataset = self.insert_dataset("Flights", [self.get_user("admin").id], main_db) + dataset = self.insert_default_dataset() prev_col_len = len(dataset.columns) + cols = [ + { + "column_name": c.column_name, + "description": c.description, + "expression": c.expression, + "type": c.type, + "advanced_data_type": c.advanced_data_type, + "verbose_name": c.verbose_name, + } + for c in dataset.columns + ] + + cols.append( + { + "column_name": "new_col", + "description": "description", + "expression": "expression", + "type": "INTEGER", + "advanced_data_type": "ADVANCED_DATA_TYPE", + "verbose_name": "New Col", + } + ) + self.login(username="admin") dataset_data = { - "sql": "SELECT * from flights ", - "columns": [ - {"column_name": "YEAR", "type": "INT", "is_dttm": False}, - {"column_name": "MONTH", "type": "INT", "is_dttm": False}, - {"column_name": "DAY", "type": "INT", "is_dttm": False}, - ], + "columns": cols, } uri = f"api/v1/dataset/{dataset.id}?override_columns=true" rv = self.put_assert_metric(uri, dataset_data, "put") + assert rv.status_code == 200 columns = db.session.query(TableColumn).filter_by(table_id=dataset.id).all()