def link_company_with_dnb(
dh_company_id,
duns_number,
modified_by,
update_descriptor='admin:link_company_with_dnb',
):
"""
Given a Data Hub company ID and a duns number, save the company with this
duns number and update it's record from D&B.
"""
company = Company.objects.get(id=dh_company_id)
if company.duns_number:
raise CompanyAlreadyDNBLinkedError(
f'Company {company.id} is already linked with duns number {company.duns_number}',
)
company.duns_number = duns_number
company.modified_by = modified_by
company.save()
sync_company_with_dnb(
company.id,
update_descriptor=update_descriptor,
retry_failures=False,
)
company.refresh_from_db()
return company
def test_sync_company_with_dnb_retries_errors(monkeypatch, error,
expect_retry):
"""
Test the sync_company_with_dnb task retries server errors.
"""
company = CompanyFactory(duns_number='123456789')
# Set up a DNBServiceError with the parametrized status code
mocked_get_company = mock.Mock()
mocked_get_company.side_effect = error
monkeypatch.setattr('datahub.dnb_api.tasks.sync.get_company',
mocked_get_company)
# Mock the task's retry method
retry_mock = mock.Mock(side_effect=Retry(exc=error))
monkeypatch.setattr('datahub.dnb_api.tasks.sync_company_with_dnb.retry',
retry_mock)
if expect_retry:
expected_exception_class = Retry
else:
expected_exception_class = DNBServiceError
with pytest.raises(expected_exception_class):
sync_company_with_dnb(company.id)
def test_sync_company_with_dnb_respects_retry_failures_flag(monkeypatch, error):
"""
Test the sync_company_with_dnb task retries server errors.
"""
company = CompanyFactory(duns_number='123456789')
# Set up a DNBServiceError with the parametrized status code
mocked_get_company = mock.Mock()
mocked_get_company.side_effect = error
monkeypatch.setattr('datahub.dnb_api.tasks.sync.get_company', mocked_get_company)
with pytest.raises(error.__class__):
sync_company_with_dnb(company.id, retry_failures=False)