def test_filename_suffix(self, filename_suffix, expected_partitions, local_csvs): """Test how specifying filename_suffix affects the available partitions and their names""" pds = IncrementalDataSet( str(local_csvs), DATASET, filename_suffix=filename_suffix ) loaded = pds.load() assert loaded.keys() == expected_partitions
def test_comparison_func(self, comparison_func, expected_partitions, local_csvs): """Test that specifying a custom function for comparing the checkpoint value to a partition id results in expected partitions being returned on load""" checkpoint_config = { "force_checkpoint": "p02/data.csv", "comparison_func": comparison_func, } pds = IncrementalDataSet(str(local_csvs), DATASET, checkpoint=checkpoint_config) assert pds.load().keys() == expected_partitions
def test_force_checkpoint_no_partitions(self, forced_checkpoint, local_csvs): """Test that forcing the checkpoint to certain values results in no partitions being returned""" pds = IncrementalDataSet(str(local_csvs), DATASET, checkpoint=forced_checkpoint) loaded = pds.load() assert loaded == {} confirm_path = local_csvs / pds.DEFAULT_CHECKPOINT_FILENAME assert not confirm_path.exists() pds.confirm() # confirming with no partitions available must have no effect assert not confirm_path.exists()
def test_force_checkpoint_no_partitions(self, forced_checkpoint, mocked_csvs_in_s3): """Test that forcing the checkpoint to certain values results in no partitions returned from S3""" pds = IncrementalDataSet( mocked_csvs_in_s3, DATASET, checkpoint=forced_checkpoint ) loaded = pds.load() assert loaded == {} assert not pds._checkpoint.exists() pds.confirm() # confirming with no partitions available must have no effect assert not pds._checkpoint.exists()
def test_force_checkpoint_checkpoint_file_exists( self, forced_checkpoint, expected_partitions, local_csvs ): """Test how forcing checkpoint value affects the available partitions if the checkpoint file exists""" IncrementalDataSet(str(local_csvs), DATASET).confirm() checkpoint = local_csvs / IncrementalDataSet.DEFAULT_CHECKPOINT_FILENAME assert checkpoint.read_text() == "p04/data.csv" pds = IncrementalDataSet(str(local_csvs), DATASET, checkpoint=forced_checkpoint) assert pds._checkpoint.exists() loaded = pds.load() assert loaded.keys() == expected_partitions
def test_force_checkpoint_no_checkpoint_file( self, forced_checkpoint, expected_partitions, local_csvs ): """Test how forcing checkpoint value affects the available partitions if the checkpoint file does not exist""" pds = IncrementalDataSet(str(local_csvs), DATASET, checkpoint=forced_checkpoint) loaded = pds.load() assert loaded.keys() == expected_partitions confirm_path = local_csvs / pds.DEFAULT_CHECKPOINT_FILENAME assert not confirm_path.exists() pds.confirm() assert confirm_path.is_file() assert confirm_path.read_text() == max(expected_partitions)
def test_force_checkpoint_no_checkpoint_file( self, forced_checkpoint, expected_partitions, mocked_csvs_in_s3 ): """Test how forcing checkpoint value affects the available partitions in S3 if the checkpoint file does not exist""" pds = IncrementalDataSet( mocked_csvs_in_s3, DATASET, checkpoint=forced_checkpoint ) loaded = pds.load() assert loaded.keys() == expected_partitions assert not pds._checkpoint.exists() pds.confirm() assert pds._checkpoint.exists() assert pds._checkpoint.load() == max(expected_partitions)
def test_credentials(self, pds_config, fs_creds, dataset_creds, checkpoint_creds): """Test correctness of credentials propagation into the dataset and checkpoint constructors""" pds = IncrementalDataSet(str(Path.cwd()), **pds_config) assert pds._credentials == fs_creds assert pds._dataset_config[CREDENTIALS_KEY] == dataset_creds assert pds._checkpoint_config[CREDENTIALS_KEY] == checkpoint_creds
def test_version_not_allowed(self, tmp_path, checkpoint_config, error_pattern): """Test that invalid checkpoint configurations raise expected errors""" with pytest.raises(DataSetError, match=re.escape(error_pattern)): IncrementalDataSet(str(tmp_path), DATASET, checkpoint=checkpoint_config)
def test_checkpoint_type(self, tmp_path, checkpoint_config, expected_checkpoint_class): """Test configuring a different checkpoint dataset type""" pds = IncrementalDataSet(str(tmp_path), DATASET, checkpoint=checkpoint_config) assert isinstance(pds._checkpoint, expected_checkpoint_class)
def test_force_checkpoint_checkpoint_file_exists(self, forced_checkpoint, expected_partitions, mocked_csvs_in_s3): """Test how forcing checkpoint value affects the available partitions in S3 if the checkpoint file exists""" # create checkpoint and assert that it exists IncrementalDataSet(mocked_csvs_in_s3, DATASET).confirm() checkpoint_path = "{}/{}".format( mocked_csvs_in_s3, IncrementalDataSet.DEFAULT_CHECKPOINT_FILENAME) checkpoint_value = TextDataSet(checkpoint_path).load() assert checkpoint_value == "p04/data.csv" pds = IncrementalDataSet(mocked_csvs_in_s3, DATASET, checkpoint=forced_checkpoint) assert pds._checkpoint.exists() loaded = pds.load() assert loaded.keys() == expected_partitions
def test_checkpoint_path(self, local_csvs, partitioned_data_pandas): """Test configuring a different checkpoint path""" checkpoint_path = local_csvs / "checkpoint_folder" / "checkpoint_file" assert not checkpoint_path.exists() IncrementalDataSet( str(local_csvs), DATASET, checkpoint={"filepath": str(checkpoint_path)} ).confirm() assert checkpoint_path.is_file() assert checkpoint_path.read_text() == max(partitioned_data_pandas)
def test_load_and_confirm(self, local_csvs, partitioned_data_pandas): """Test the standard flow for loading, confirming and reloading an IncrementalDataSet""" pds = IncrementalDataSet(str(local_csvs), DATASET) loaded = pds.load() assert loaded.keys() == partitioned_data_pandas.keys() for partition_id, data in loaded.items(): assert_frame_equal(data, partitioned_data_pandas[partition_id]) checkpoint_path = local_csvs / pds.DEFAULT_CHECKPOINT_FILENAME assert not checkpoint_path.exists() pds.confirm() assert checkpoint_path.is_file() assert checkpoint_path.read_text() == pds._read_checkpoint() == "p04/data.csv" reloaded = pds.load() assert reloaded.keys() == loaded.keys() pds.release() reloaded_after_release = pds.load() assert reloaded_after_release == {}
def test_load_and_confirm(self, mocked_csvs_in_s3, partitioned_data_pandas): """Test the standard flow for loading, confirming and reloading a IncrementalDataSet in S3""" pds = IncrementalDataSet(mocked_csvs_in_s3, DATASET) assert pds._checkpoint._protocol == "s3" loaded = pds.load() assert loaded.keys() == partitioned_data_pandas.keys() for partition_id, data in loaded.items(): assert_frame_equal(data, partitioned_data_pandas[partition_id]) assert not pds._checkpoint.exists() assert pds._read_checkpoint() is None pds.confirm() assert pds._checkpoint.exists() assert pds._read_checkpoint() == max(partitioned_data_pandas)
def test_load_and_confirm_s3a(self, mocked_csvs_in_s3, partitioned_data_pandas, mocker): s3a_path = f"s3a://{mocked_csvs_in_s3.split('://', 1)[1]}" pds = IncrementalDataSet(s3a_path, DATASET) assert pds._protocol == "s3a" assert pds._checkpoint._protocol == "s3" mocked_ds = mocker.patch.object(pds, "_dataset_type") mocked_ds.__name__ = "mocked" loaded = pds.load() assert loaded.keys() == partitioned_data_pandas.keys() assert not pds._checkpoint.exists() assert pds._read_checkpoint() is None pds.confirm() assert pds._checkpoint.exists() assert pds._read_checkpoint() == max(partitioned_data_pandas)
def test_save(self, local_csvs): """Test saving a new partition into an IncrementalDataSet""" df = pd.DataFrame({"dummy": [1, 2, 3]}) new_partition_key = "p05/data.csv" new_partition_path = local_csvs / new_partition_key pds = IncrementalDataSet(str(local_csvs), DATASET) assert not new_partition_path.exists() assert new_partition_key not in pds.load() pds.save({new_partition_key: df}) assert new_partition_path.exists() loaded = pds.load() assert_frame_equal(loaded[new_partition_key], df)