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)
