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]