def test_init_dataset_with_alt_project_and_location( self, get_dataset_mock): aiplatform.init(project=_TEST_PROJECT) datasets._Dataset(dataset_name=_TEST_NAME, project=_TEST_ALT_PROJECT, location=_TEST_LOCATION) get_dataset_mock.assert_called_once_with(name=_TEST_NAME)
def test_init_dataset_with_project_and_alt_location(self): aiplatform.init(project=_TEST_PROJECT) with pytest.raises(RuntimeError): datasets._Dataset( dataset_name=_TEST_NAME, project=_TEST_PROJECT, location=_TEST_ALT_LOCATION, )
def test_export_data(self, export_data_mock): aiplatform.init(project=_TEST_PROJECT) my_dataset = datasets._Dataset(dataset_name=_TEST_NAME) my_dataset.export_data(output_dir=_TEST_OUTPUT_DIR) expected_export_config = gca_dataset.ExportDataConfig( gcs_destination=gca_io.GcsDestination( output_uri_prefix=_TEST_OUTPUT_DIR)) export_data_mock.assert_called_once_with( name=_TEST_NAME, export_config=expected_export_config)
def test_import_data(self, import_data_mock, sync): aiplatform.init(project=_TEST_PROJECT) my_dataset = datasets._Dataset(dataset_name=_TEST_NAME) my_dataset.import_data( gcs_source=_TEST_SOURCE_URI_GCS, import_schema_uri=_TEST_IMPORT_SCHEMA_URI, data_item_labels=_TEST_DATA_LABEL_ITEMS, sync=sync, ) if not sync: my_dataset.wait() expected_import_config = gca_dataset.ImportDataConfig( gcs_source=gca_io.GcsSource(uris=[_TEST_SOURCE_URI_GCS]), import_schema_uri=_TEST_IMPORT_SCHEMA_URI, data_item_labels=_TEST_DATA_LABEL_ITEMS, ) import_data_mock.assert_called_once_with( name=_TEST_NAME, import_configs=[expected_import_config])
def test_init_dataset_with_invalid_name(self): with pytest.raises(ValueError): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) datasets._Dataset(dataset_name=_TEST_INVALID_NAME)
def test_init_dataset_with_id_only_without_project_or_location(self): with pytest.raises(GoogleAuthError): datasets._Dataset( dataset_name=_TEST_ID, credentials=auth_credentials.AnonymousCredentials(), )
def test_init_dataset_with_id_only(self, get_dataset_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) datasets._Dataset(dataset_name=_TEST_ID) get_dataset_mock.assert_called_once_with(name=_TEST_NAME)
def test_init_dataset(self, get_dataset_mock): aiplatform.init(project=_TEST_PROJECT) datasets._Dataset(dataset_name=_TEST_NAME) get_dataset_mock.assert_called_once_with(name=_TEST_NAME)