address concerns

apache · Jul 30, 2022 · 12963aa · 12963aa
1 parent 9e1a585
commit 12963aa
Showing 1 changed file with 26 additions and 7 deletions.
diff --git 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()