def test_ingestion_manager_run_success():
df = pd.DataFrame({"float": pd.Series([2.0], dtype="float64")})
manager = IngestionManagerPandas(
feature_group_name="MyGroup",
sagemaker_session=sagemaker_session_mock,
data_frame=df,
max_workers=10,
)
manager.run()
def test_ingestion_manager_run_multi_process_with_multi_thread_success(
fs_runtime_client_config_mock, ):
df = pd.DataFrame({"float": pd.Series([2.0], dtype="float64")})
manager = IngestionManagerPandas(
feature_group_name="MyGroup",
sagemaker_fs_runtime_client_config=fs_runtime_client_config_mock,
max_workers=2,
max_processes=2,
)
manager.run(df)
def test_ingestion_manager_run_success():
df = pd.DataFrame({"float": pd.Series([2.0], dtype="float64")})
manager = IngestionManagerPandas(
feature_group_name="MyGroup",
sagemaker_fs_runtime_client_config=fs_runtime_client_config_mock,
max_workers=10,
)
manager.run(df)
manager._run_multi_process.assert_called_once_with(data_frame=df, wait=True, timeout=None)
def test_ingestion_manager_run_failure():
df = pd.DataFrame({"float": pd.Series([2.0], dtype="float64")})
manager = IngestionManagerPandas(
feature_group_name="MyGroup",
sagemaker_session=sagemaker_session_mock,
data_frame=df,
max_workers=10,
)
with pytest.raises(RuntimeError) as error:
manager.run()
assert "Failed to ingest some data into FeatureGroup MyGroup" in str(error)
def test_ingestion_manager_with_profile_name_run_failure():
df = pd.DataFrame({"float": pd.Series([2.0], dtype="float64")})
manager = IngestionManagerPandas(
feature_group_name="MyGroup",
sagemaker_fs_runtime_client_config=fs_runtime_client_config_mock,
max_workers=1,
profile_name="non_exist",
)
try:
manager.run(df)
except Exception as e:
assert "The config profile (non_exist) could not be found" in str(e)
def test_ingestion_manager_run_failure():
df = pd.DataFrame({"float": pd.Series([2.0], dtype="float64")})
manager = IngestionManagerPandas(
feature_group_name="MyGroup",
sagemaker_fs_runtime_client_config=fs_runtime_client_config_mock,
max_workers=1,
)
with pytest.raises(IngestionError) as error:
manager.run(df)
assert "Failed to ingest some data into FeatureGroup MyGroup" in str(error)
assert error.value.failed_rows == [1]
assert manager.failed_rows == [1]