def test_dataset_update(capsys): dataset = ensure_dataset_ready() automl_tables_dataset.update_dataset(PROJECT, REGION, dataset.display_name, target_column_spec_name='Deposit', weight_column_spec_name='Balance') out, _ = capsys.readouterr() assert "Target column updated." in out assert "Weight column updated." in out
def ensure_dataset_ready(): dataset = None name = STATIC_DATASET try: dataset = automl_tables_dataset.get_dataset(PROJECT, REGION, name) except exceptions.NotFound: dataset = automl_tables_dataset.create_dataset(PROJECT, REGION, name) if dataset.example_count is None or dataset.example_count == 0: automl_tables_dataset.import_data(PROJECT, REGION, name, GCS_DATASET) dataset = automl_tables_dataset.get_dataset(PROJECT, REGION, name) automl_tables_dataset.update_dataset( PROJECT, REGION, dataset.display_name, target_column_spec_name="Deposit", ) return dataset