def test_run__ignore_error_throttling(self, dml_mock, query_mock):
task = _make_task(
DeleteData,
{
"options": {
"objects": "Contact",
"ignore_row_errors": "true",
"hardDelete": "true",
}
},
)
query_mock.return_value.get_results.return_value = iter(
["001000000000000", "001000000000001"] * 15)
query_mock.return_value.job_result = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 2, 0)
dml_mock.return_value.get_results.return_value = iter([
DataOperationResult("001000000000000", True, None),
DataOperationResult("001000000000001", False, None),
] * 15)
dml_mock.return_value.job_result = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 2, 0)
with mock.patch.object(task.logger, "warning") as warning:
task()
assert len(warning.mock_calls) == task.row_warning_limit + 1 == 11
assert "warnings suppressed" in str(warning.mock_calls[-1])
def test_run(self, dml_mock, query_mock):
task = _make_task(DeleteData, {"options": {"objects": "Contact"}})
query_mock.return_value.get_results.return_value = iter(
["001000000000000", "001000000000001"])
query_mock.return_value.job_result = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 2, 0)
dml_mock.return_value.get_results.return_value = iter([
DataOperationResult("001000000000000", True, None),
DataOperationResult("001000000000001", True, None),
])
dml_mock.return_value.job_result = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 0, 0)
task()
query_mock.assert_called_once_with(
sobject="Contact",
api_options={},
context=task,
query="SELECT Id FROM Contact",
)
query_mock.return_value.query.assert_called_once()
query_mock.return_value.get_results.assert_called_once()
dml_mock.assert_called_once_with(
sobject="Contact",
operation=DataOperationType.DELETE,
api_options={},
context=task,
fields=["Id"],
)
dml_mock.return_value.start.assert_called_once()
dml_mock.return_value.end.assert_called_once()
dml_mock.return_value.load_records.assert_called_once()
dml_mock.return_value.get_results.assert_called_once()
def test_run_task__after_steps_failure(self):
task = _make_task(
LoadData,
{
"options": {
"database_url": "sqlite://",
"mapping": "mapping.yml"
}
},
)
task._init_db = mock.Mock()
task._init_mapping = mock.Mock()
task._expand_mapping = mock.Mock()
task.mapping = {}
task.mapping["Insert Households"] = 1
task.mapping["Insert Contacts"] = 2
households_steps = {}
households_steps["four"] = 4
households_steps["five"] = 5
task.after_steps = {
"Insert Contacts": {
"three": 3
},
"Insert Households": households_steps,
}
task._execute_step = mock.Mock(side_effect=[
DataOperationJobResult(DataOperationStatus.SUCCESS, [], 0, 0),
DataOperationJobResult(DataOperationStatus.JOB_FAILURE, [], 0, 0),
])
with self.assertRaises(BulkDataException):
task()
def test_run__ignore_error(self, dml_mock, query_mock):
mock_describe_calls()
task = _make_task(
DeleteData,
{
"options": {
"objects": "Contact",
"ignore_row_errors": "true",
"hardDelete": "true",
}
},
)
query_mock.return_value.get_results.return_value = iter(
["001000000000000", "001000000000001"])
query_mock.return_value.job_result = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 2, 0)
dml_mock.return_value.get_results.return_value = iter([
DataOperationResult("001000000000000", True, None),
DataOperationResult("001000000000001", False, None),
])
dml_mock.return_value.job_result = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 2, 0)
with mock.patch.object(task.logger, "warning") as warning:
task()
assert len(warning.mock_calls) == 1
query_mock.assert_called_once_with(
sobject="Contact",
api_options={},
context=task,
query="SELECT Id FROM Contact",
api=DataApi.SMART,
fields=["Id"],
)
query_mock.return_value.query.assert_called_once()
query_mock.return_value.get_results.assert_called_once()
dml_mock.assert_called_once_with(
sobject="Contact",
operation=DataOperationType.HARD_DELETE,
api_options={},
context=task,
fields=["Id"],
api=DataApi.SMART,
volume=2,
)
dml_mock.return_value.start.assert_called_once()
dml_mock.return_value.end.assert_called_once()
dml_mock.return_value.load_records.assert_called_once()
dml_mock.return_value.get_results.assert_called_once()
def test_query(self):
context = mock.Mock()
context.sf.query.return_value = {
"totalSize": 2,
"done": True,
"records": [
{
"Id": "003000000000001",
"LastName": "Narvaez",
"Email": "*****@*****.**",
},
{"Id": "003000000000002", "LastName": "De Vries", "Email": None},
],
}
query_op = RestApiQueryOperation(
sobject="Contact",
fields=["Id", "LastName", "Email"],
api_options={},
context=context,
query="SELECT Id, LastName, Email FROM Contact",
)
query_op.query()
assert query_op.job_result == DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 2, 0
)
assert list(query_op.get_results()) == [
["003000000000001", "Narvaez", "*****@*****.**"],
["003000000000002", "De Vries", ""],
]
def test_query(self):
context = mock.Mock()
query = BulkApiQueryOperation(
sobject="Contact",
api_options={},
context=context,
query="SELECT Id FROM Contact",
)
query._wait_for_job = mock.Mock()
query._wait_for_job.return_value = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 0, 0)
query.query()
assert query.job_result.status is DataOperationStatus.SUCCESS
context.bulk.create_query_job.assert_called_once_with(
"Contact", contentType="CSV")
context.bulk.query.assert_called_once_with(
context.bulk.create_query_job.return_value,
"SELECT Id FROM Contact")
query._wait_for_job.assert_called_once_with(
context.bulk.create_query_job.return_value)
context.bulk.close_job.assert_called_once_with(
context.bulk.create_query_job.return_value)
def test_get_results__no_results(self, download_mock):
context = mock.Mock()
context.bulk.endpoint = "https://test"
context.bulk.create_query_job.return_value = "JOB"
context.bulk.query.return_value = "BATCH"
context.bulk.get_query_batch_result_ids.return_value = ["RESULT"]
download_mock.return_value = io.StringIO(
"Records not found for this query")
query = BulkApiQueryOperation(
sobject="Contact",
api_options={},
context=context,
query="SELECT Id FROM Contact",
)
query._wait_for_job = mock.Mock()
query._wait_for_job.return_value = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 0, 0)
query.query()
results = list(query.get_results())
context.bulk.get_query_batch_result_ids.assert_called_once_with(
"BATCH", job_id="JOB")
download_mock.assert_called_once_with(
"https://test/job/JOB/batch/BATCH/result/RESULT", context.bulk)
assert list(results) == []
def test_contextmanager(self):
context = mock.Mock()
context.bulk.create_job.return_value = "JOB"
step = BulkApiDmlOperation(
sobject="Contact",
operation=DataOperationType.INSERT,
api_options={},
context=context,
fields=["LastName"],
)
step._wait_for_job = mock.Mock()
step._wait_for_job.return_value = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 0, 0)
step.job_id = "JOB"
with step:
pass
context.bulk.create_job.assert_called_once_with("Contact",
"insert",
contentType="CSV",
concurrency="Parallel")
assert step.job_id == "JOB"
context.bulk.close_job.assert_called_once_with("JOB")
step._wait_for_job.assert_called_once_with("JOB")
assert step.job_result.status is DataOperationStatus.SUCCESS
def test_run_task__after_steps(self):
task = _make_task(
LoadData,
{
"options": {
"database_url": "sqlite://",
"mapping": "mapping.yml"
}
},
)
task._init_db = mock.Mock()
task._init_mapping = mock.Mock()
task._expand_mapping = mock.Mock()
task.mapping = {}
task.mapping["Insert Households"] = 1
task.mapping["Insert Contacts"] = 2
households_steps = {}
households_steps["four"] = 4
households_steps["five"] = 5
task.after_steps = {
"Insert Contacts": {
"three": 3
},
"Insert Households": households_steps,
}
task._execute_step = mock.Mock(return_value=DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 0, 0))
task()
task._execute_step.assert_has_calls([
mock.call(1),
mock.call(4),
mock.call(5),
mock.call(2),
mock.call(3)
])
def test_end_to_end(self, download_mock):
context = mock.Mock()
context.bulk.endpoint = "https://test"
context.bulk.create_job.return_value = "JOB"
context.bulk.post_batch.side_effect = ["BATCH1", "BATCH2"]
download_mock.return_value = io.StringIO("""id,success,created,error
003000000000001,true,true,
003000000000002,true,true,
003000000000003,false,false,error""")
step = BulkApiDmlOperation(
sobject="Contact",
operation=DataOperationType.INSERT,
api_options={},
context=context,
fields=["LastName"],
)
step._wait_for_job = mock.Mock()
step._wait_for_job.return_value = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 0, 0)
step.start()
step.load_records(iter([["Test"], ["Test2"], ["Test3"]]))
step.end()
assert step.job_result.status is DataOperationStatus.SUCCESS
results = step.get_results()
assert list(results) == [
DataOperationResult("003000000000001", True, None),
DataOperationResult("003000000000002", True, None),
DataOperationResult(None, False, "error"),
]
def test_insert_dml_operation(self):
mock_describe_calls()
task = _make_task(
LoadData,
{
"options": {
"database_url": "sqlite:///test.db",
"mapping": "mapping.yml",
}
},
)
task.project_config.project__package__api_version = "48.0"
task._init_task()
responses.add(
responses.POST,
url="https://example.com/services/data/v48.0/composite/sobjects",
json=[
{
"id": "003000000000001",
"success": True
},
{
"id": "003000000000002",
"success": True
},
],
status=200,
)
responses.add(
responses.POST,
url="https://example.com/services/data/v48.0/composite/sobjects",
json=[{
"id": "003000000000003",
"success": True
}],
status=200,
)
recs = [["Fred", "Narvaez"], [None, "De Vries"], ["Hiroko", "Aito"]]
dml_op = RestApiDmlOperation(
sobject="Contact",
operation=DataOperationType.INSERT,
api_options={"batch_size": 2},
context=task,
fields=["FirstName", "LastName"],
)
dml_op.start()
dml_op.load_records(iter(recs))
dml_op.end()
assert dml_op.job_result == DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 3, 0)
assert list(dml_op.get_results()) == [
DataOperationResult("003000000000001", True, ""),
DataOperationResult("003000000000002", True, ""),
DataOperationResult("003000000000003", True, ""),
]
def _job_state_from_batches(self, job_id):
return DataOperationJobResult(
DataOperationStatus.SUCCESS,
[],
10,
0,
)
def test_run__job_error_query(self, dml_mock, query_mock):
task = _make_task(DeleteData, {"options": {"objects": "Contact"}})
query_mock.return_value.get_results.return_value = iter(
["001000000000000", "001000000000001"])
query_mock.return_value.job_result = DataOperationJobResult(
DataOperationStatus.JOB_FAILURE, [], 0, 0)
with self.assertRaises(BulkDataException):
task()
def test_wait_for_job(self, sleep_patch):
mixin = BulkJobMixin()
mixin.bulk = mock.Mock()
mixin.bulk.job_status.return_value = {
"numberBatchesCompleted": 1,
"numberBatchesTotal": 1,
}
mixin._job_state_from_batches = mock.Mock(side_effect=[
DataOperationJobResult(DataOperationStatus.IN_PROGRESS, [], 0, 0),
DataOperationJobResult(DataOperationStatus.SUCCESS, [], 0, 0),
])
mixin.logger = mock.Mock()
result = mixin._wait_for_job("750000000000000")
mixin._job_state_from_batches.assert_has_calls(
[mock.call("750000000000000"),
mock.call("750000000000000")])
assert result.status is DataOperationStatus.SUCCESS
def test_run_query__failure(self, query_op_mock):
task = _make_task(
ExtractData, {"options": {"database_url": "sqlite:///", "mapping": ""}}
)
query_op_mock.return_value.job_result = DataOperationJobResult(
DataOperationStatus.JOB_FAILURE, [], 1, 0
)
with self.assertRaises(BulkDataException):
task._run_query("SELECT Id FROM Contact", MappingStep(sf_object="Contact"))
def test_run__job_error_delete(self, dml_mock, query_mock):
task = _make_task(DeleteData, {"options": {"objects": "Contact"}})
query_mock.return_value.get_results.return_value = iter(
["001000000000000", "001000000000001"])
query_mock.return_value.job_result = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 2, 0)
dml_mock.return_value.get_results.return_value = iter([
DataOperationResult("001000000000000", True, None),
DataOperationResult("001000000000001", False, None),
])
with self.assertRaises(BulkDataException):
task()
def test_parse_job_state(self):
mixin = BulkJobMixin()
mixin.bulk = mock.Mock()
mixin.bulk.jobNS = "http://ns"
assert (mixin._parse_job_state(
BULK_BATCH_RESPONSE.format(
**{
"first_state": "Not Processed",
"first_message": "Test",
"second_state": "Completed",
"second_message": "",
})) == DataOperationJobResult(DataOperationStatus.ABORTED, [],
0, 0))
assert (mixin._parse_job_state(
BULK_BATCH_RESPONSE.format(
**{
"first_state": "InProgress",
"first_message": "Test",
"second_state": "Completed",
"second_message": "",
})) == DataOperationJobResult(DataOperationStatus.IN_PROGRESS,
[], 0, 0))
assert mixin._parse_job_state(
BULK_BATCH_RESPONSE.format(
**{
"first_state": "Failed",
"first_message": "Bad",
"second_state": "Failed",
"second_message": "Worse",
})) == DataOperationJobResult(DataOperationStatus.JOB_FAILURE,
["Bad", "Worse"], 0, 0)
assert (mixin._parse_job_state(
BULK_BATCH_RESPONSE.format(
**{
"first_state": "Completed",
"first_message": "Test",
"second_state": "Completed",
"second_message": "",
})) == DataOperationJobResult(DataOperationStatus.SUCCESS, [],
0, 0))
assert mixin._parse_job_state(
'<root xmlns="http://ns">'
" <batch><state>Completed</state></batch>"
" <numberRecordsFailed>200</numberRecordsFailed>"
"</root>") == DataOperationJobResult(
DataOperationStatus.ROW_FAILURE, [], 0, 200)
assert mixin._parse_job_state(
'<root xmlns="http://ns">'
" <batch><state>Completed</state></batch>"
" <numberRecordsFailed>200</numberRecordsFailed>"
" <numberRecordsProcessed>10</numberRecordsProcessed>"
"</root>") == DataOperationJobResult(
DataOperationStatus.ROW_FAILURE, [], 10, 200)
def test_query__contextmanager(self):
context = mock.Mock()
query = BulkApiQueryOperation(
sobject="Contact",
api_options={},
context=context,
query="SELECT Id FROM Contact",
)
query._wait_for_job = mock.Mock()
query._wait_for_job.return_value = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 0, 0)
with query:
assert query.job_result.status is DataOperationStatus.SUCCESS
def test_run__query_fails(self, dml_mock, query_mock):
mock_describe_calls()
task = _make_task(
DeleteData,
{"options": {
"objects": "Contact",
"where": "Id != null"
}})
query_mock.return_value.get_results.return_value = iter(
["001000000000000", "001000000000001"])
query_mock.return_value.job_result = DataOperationJobResult(
DataOperationStatus.JOB_FAILURE, [], 0, 0)
with self.assertRaises(BulkDataException):
task()
def test_query__failure(self):
context = mock.Mock()
query = BulkApiQueryOperation(
sobject="Contact",
api_options={},
context=context,
query="SELECT Id FROM Contact",
)
query._wait_for_job = mock.Mock()
query._wait_for_job.return_value = DataOperationJobResult(
DataOperationStatus.JOB_FAILURE, [], 0, 0)
query.query()
assert query.job_result.status is DataOperationStatus.JOB_FAILURE
def test_wait_for_job__logs_state_messages(self):
mixin = BulkJobMixin()
mixin.bulk = mock.Mock()
mixin.bulk.job_status.return_value = {
"numberBatchesCompleted": 1,
"numberBatchesTotal": 1,
}
mixin._job_state_from_batches = mock.Mock(
return_value=DataOperationJobResult(
DataOperationStatus.JOB_FAILURE, ["Test1", "Test2"], 0, 0))
mixin.logger = mock.Mock()
mixin._wait_for_job("750000000000000")
mixin.logger.error.assert_any_call("Batch failure message: Test1")
mixin.logger.error.assert_any_call("Batch failure message: Test2")
def test_process_job_results__insert_rows_fail(self):
task = _make_task(
LoadData,
{
"options": {
"database_url": "sqlite://",
"mapping": "mapping.yml",
"ignore_row_errors": True,
}
},
)
task.session = mock.Mock()
task._initialize_id_table = mock.Mock()
task._sql_bulk_insert_from_records = mock.Mock()
task.bulk = mock.Mock()
task.sf = mock.Mock()
task.logger = mock.Mock()
local_ids = ["1", "2", "3", "4"]
step = MockBulkApiDmlOperation(
sobject="Contact",
operation=DataOperationType.INSERT,
api_options={},
context=task,
fields=[],
)
step.job_result = DataOperationJobResult(
DataOperationStatus.ROW_FAILURE, [], 4, 4)
step.end = mock.Mock()
step.results = [
DataOperationResult("001111111111111", False, None),
DataOperationResult("001111111111112", False, None),
DataOperationResult("001111111111113", False, None),
DataOperationResult("001111111111114", False, None),
]
mapping = {"table": "Account", "action": "insert"}
task._process_job_results(mapping, step, local_ids)
task.session.connection.assert_called_once()
task._initialize_id_table.assert_called_once_with(mapping, True)
task._sql_bulk_insert_from_records.assert_not_called()
task.session.commit.assert_called_once()
assert len(task.logger.mock_calls) == 4
def test_wait_for_job__failed(self):
mixin = BulkJobMixin()
mixin.bulk = mock.Mock()
mixin.bulk.job_status.return_value = {
"numberBatchesCompleted": 1,
"numberBatchesTotal": 1,
}
mixin._job_state_from_batches = mock.Mock(
return_value=DataOperationJobResult(
DataOperationStatus.JOB_FAILURE, ["Test1", "Test2"], 0, 0))
mixin.logger = mock.Mock()
result = mixin._wait_for_job("750000000000000")
mixin._job_state_from_batches.assert_called_once_with(
"750000000000000")
assert result.status is DataOperationStatus.JOB_FAILURE
def test_run__no_results(self, dml_mock, query_mock):
task = _make_task(DeleteData, {"options": {"objects": "Contact"}})
query_mock.return_value.get_results.return_value = iter([])
query_mock.return_value.job_result = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 0, 0)
task()
query_mock.assert_called_once_with(
sobject="Contact",
api_options={},
context=task,
query="SELECT Id FROM Contact",
)
query_mock.return_value.query.assert_called_once()
query_mock.return_value.get_results.assert_not_called()
dml_mock.assert_not_called()
def test_run_task__exception_failure(self):
task = _make_task(
LoadData,
{
"options": {
"database_url": "sqlite://",
"mapping": "mapping.yml"
}
},
)
task._init_db = mock.Mock()
task._init_mapping = mock.Mock()
task._execute_step = mock.Mock(return_value=DataOperationJobResult(
DataOperationStatus.JOB_FAILURE, [], 0, 0))
task.mapping = {"Test": {"test": "test"}}
with self.assertRaises(BulkDataException):
task()
def test_run_query__no_results(self, step_mock):
task = _make_task(
ExtractData,
{"options": {
"database_url": "sqlite:///",
"mapping": ""
}})
task._import_results = mock.Mock()
step_mock.return_value.job_result = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 0, 0)
task._run_query("SELECT Id FROM Contact", {"sf_object": "Contact"})
step_mock.assert_called_once_with(
sobject="Contact",
api_options={},
context=task,
query="SELECT Id FROM Contact",
)
step_mock.return_value.query.assert_called_once_with()
task._import_results.assert_not_called()
def test_end__failed(self):
context = mock.Mock()
context.bulk.create_job.return_value = "JOB"
step = BulkApiDmlOperation(
sobject="Contact",
operation=DataOperationType.INSERT,
api_options={},
context=context,
fields=["LastName"],
)
step._wait_for_job = mock.Mock()
step._wait_for_job.return_value = DataOperationJobResult(
DataOperationStatus.JOB_FAILURE, [], 0, 0)
step.job_id = "JOB"
step.end()
context.bulk.close_job.assert_called_once_with("JOB")
step._wait_for_job.assert_called_once_with("JOB")
assert step.job_result.status is DataOperationStatus.JOB_FAILURE
def test_run_query(self, query_op_mock):
task = _make_task(
ExtractData, {"options": {"database_url": "sqlite:///", "mapping": ""}}
)
task._import_results = mock.Mock()
query_op_mock.return_value.job_result = DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 1, 0
)
task._run_query("SELECT Id FROM Contact", MappingStep(sf_object="Contact"))
query_op_mock.assert_called_once_with(
sobject="Contact",
fields=["Id"],
api=DataApi.SMART,
api_options={},
context=task,
query="SELECT Id FROM Contact",
)
query_op_mock.return_value.query.assert_called_once_with()
task._import_results.assert_called_once_with(
MappingStep(sf_object="Contact"), query_op_mock.return_value
)
def test_run_task__start_step(self):
task = _make_task(
LoadData,
{
"options": {
"database_url": "sqlite://",
"mapping": "mapping.yml",
"start_step": "Insert Contacts",
}
},
)
task._init_db = mock.Mock()
task._init_mapping = mock.Mock()
task.mapping = {}
task.mapping["Insert Households"] = MappingStep(sf_object="one",
fields={})
task.mapping["Insert Contacts"] = MappingStep(sf_object="two",
fields={})
task.after_steps = {}
task._execute_step = mock.Mock(return_value=DataOperationJobResult(
DataOperationStatus.SUCCESS, [], 0, 0))
task()
task._execute_step.assert_called_once_with(
MappingStep(sf_object="two", fields={}))
def end(self):
records_processed = len(self.results)
self.job_result = DataOperationJobResult(DataOperationStatus.SUCCESS,
[], records_processed, 0)
