Merge pull request #2528 from alphagov/delete-count

Fix delete count

app/dao/notifications_dao.py

@@ -310,9 +310,7 @@ def delete_notifications_older_than_retention_by_type(notification_type, qry_lim
 
         current_app.logger.info(
             "Deleting {} notifications for service id: {}".format(notification_type, f.service_id))
-        deleted += _delete_notifications(
-            deleted, notification_type, days_of_retention, f.service_id, qry_limit
-        )
+        deleted += _delete_notifications(notification_type, days_of_retention, f.service_id, qry_limit)
 
     current_app.logger.info(
         'Deleting {} notifications for services without flexible data retention'.format(notification_type))
@@ -327,16 +325,14 @@ def delete_notifications_older_than_retention_by_type(notification_type, qry_lim
                 notification_type, service_id, seven_days_ago, qry_limit
             )
         insert_update_notification_history(notification_type, seven_days_ago, service_id)
-        deleted += _delete_notifications(
-            deleted, notification_type, seven_days_ago, service_id, qry_limit
-        )
+        deleted += _delete_notifications(notification_type, seven_days_ago, service_id, qry_limit)
 
     current_app.logger.info('Finished deleting {} notifications'.format(notification_type))
 
     return deleted
 
 
-def _delete_notifications(deleted, notification_type, date_to_delete_from, service_id, query_limit):
+def _delete_notifications(notification_type, date_to_delete_from, service_id, query_limit):
     subquery = db.session.query(
         Notification.id
     ).join(NotificationHistory, NotificationHistory.id == Notification.id).filter(
@@ -345,7 +341,7 @@ def _delete_notifications(deleted, notification_type, date_to_delete_from, servi
         Notification.created_at < date_to_delete_from,
     ).limit(query_limit).subquery()
 
-    deleted += _delete_for_query(subquery)
+    deleted = _delete_for_query(subquery)
 
     subquery_for_test_keys = db.session.query(
         Notification.id

tests/app/celery/test_nightly_tasks.py

@@ -105,8 +105,8 @@ def test_will_remove_csv_files_for_jobs_older_than_retention_period(
     mocker.patch('app.celery.nightly_tasks.s3.remove_job_from_s3')
     service_1 = create_service(service_name='service 1')
     service_2 = create_service(service_name='service 2')
-    create_service_data_retention(service_id=service_1.id, notification_type=SMS_TYPE, days_of_retention=3)
-    create_service_data_retention(service_id=service_2.id, notification_type=EMAIL_TYPE, days_of_retention=30)
+    create_service_data_retention(service=service_1, notification_type=SMS_TYPE, days_of_retention=3)
+    create_service_data_retention(service=service_2, notification_type=EMAIL_TYPE, days_of_retention=30)
     sms_template_service_1 = create_template(service=service_1)
     email_template_service_1 = create_template(service=service_1, template_type='email')
 

tests/app/dao/notification_dao/test_notification_dao_delete_notifications.py

@@ -20,6 +20,40 @@ from tests.app.db import (
 )
 
 
+def create_test_data(notification_type, sample_service, days_of_retention=3):
+    service_with_default_data_retention = create_service(service_name='default data retention')
+    email_template, letter_template, sms_template = _create_templates(sample_service)
+    default_email_template, default_letter_template, default_sms_template = _create_templates(
+        service_with_default_data_retention)
+    create_notification(template=email_template, status='delivered')
+    create_notification(template=sms_template, status='permanent-failure')
+    create_notification(template=letter_template, status='temporary-failure',
+                        reference='LETTER_REF', sent_at=datetime.utcnow())
+    create_notification(template=email_template, status='delivered',
+                        created_at=datetime.utcnow() - timedelta(days=4))
+    create_notification(template=sms_template, status='permanent-failure',
+                        created_at=datetime.utcnow() - timedelta(days=4))
+    create_notification(template=letter_template, status='temporary-failure',
+                        reference='LETTER_REF', sent_at=datetime.utcnow(),
+                        created_at=datetime.utcnow() - timedelta(days=4))
+    create_notification(template=default_email_template, status='delivered',
+                        created_at=datetime.utcnow() - timedelta(days=8))
+    create_notification(template=default_sms_template, status='permanent-failure',
+                        created_at=datetime.utcnow() - timedelta(days=8))
+    create_notification(template=default_letter_template, status='temporary-failure',
+                        reference='LETTER_REF', sent_at=datetime.utcnow(),
+                        created_at=datetime.utcnow() - timedelta(days=8))
+    create_service_data_retention(service=sample_service, notification_type=notification_type,
+                                  days_of_retention=days_of_retention)
+
+
+def _create_templates(sample_service):
+    email_template = create_template(service=sample_service, template_type='email')
+    sms_template = create_template(service=sample_service)
+    letter_template = create_template(service=sample_service, template_type='letter')
+    return email_template, letter_template, sms_template
+
+
 @pytest.mark.parametrize('month, delete_run_time',
                          [(4, '2016-04-10 23:40'), (1, '2016-01-11 00:40')])
 @pytest.mark.parametrize(
@@ -39,7 +73,6 @@ def test_should_delete_notifications_by_type_after_seven_days(
         expected_letter_count
 ):
     mocker.patch("app.dao.notifications_dao.get_s3_bucket_objects")
-    assert len(Notification.query.all()) == 0
     email_template, letter_template, sms_template = _create_templates(sample_service)
     # create one notification a day between 1st and 10th from 11:00 to 19:00 of each type
     for i in range(1, 11):
@@ -48,11 +81,12 @@ def test_should_delete_notifications_by_type_after_seven_days(
             create_notification(template=email_template, created_at=datetime.utcnow(), status="permanent-failure")
             create_notification(template=sms_template, created_at=datetime.utcnow(), status="delivered")
             create_notification(template=letter_template, created_at=datetime.utcnow(), status="temporary-failure")
-    all_notifications = Notification.query.all()
-    assert len(all_notifications) == 30
+    assert Notification.query.count() == 30
+
     # Records from before 3rd should be deleted
     with freeze_time(delete_run_time):
         delete_notifications_older_than_retention_by_type(notification_type)
+
     remaining_sms_notifications = Notification.query.filter_by(notification_type='sms').all()
     remaining_letter_notifications = Notification.query.filter_by(notification_type='letter').all()
     remaining_email_notifications = Notification.query.filter_by(notification_type='email').all()
@@ -70,9 +104,8 @@ def test_should_delete_notifications_by_type_after_seven_days(
         assert notification.created_at.date() >= date(2016, month, 3)
 
 
-@pytest.mark.parametrize('notification_type', ['sms', 'email', 'letter'])
 @freeze_time("2016-01-10 12:00:00.000000")
-def test_should_not_delete_notification_history(sample_service, notification_type, mocker):
+def test_should_not_delete_notification_history(sample_service, mocker):
     mocker.patch("app.dao.notifications_dao.get_s3_bucket_objects")
     with freeze_time('2016-01-01 12:00'):
         email_template, letter_template, sms_template = _create_templates(sample_service)
@@ -80,8 +113,7 @@ def test_should_not_delete_notification_history(sample_service, notification_typ
         create_notification(template=sms_template, status='permanent-failure')
         create_notification(template=letter_template, status='permanent-failure')
     assert Notification.query.count() == 3
-    assert NotificationHistory.query.count() == 0
-    delete_notifications_older_than_retention_by_type(notification_type)
+    delete_notifications_older_than_retention_by_type('sms')
     assert Notification.query.count() == 2
     assert NotificationHistory.query.count() == 1
 
@@ -90,10 +122,10 @@ def test_should_not_delete_notification_history(sample_service, notification_typ
 def test_delete_notifications_for_days_of_retention(sample_service, notification_type, mocker):
     mock_get_s3 = mocker.patch("app.dao.notifications_dao.get_s3_bucket_objects")
     create_test_data(notification_type, sample_service)
-    assert len(Notification.query.all()) == 9
+    assert Notification.query.count() == 9
     delete_notifications_older_than_retention_by_type(notification_type)
-    assert len(Notification.query.all()) == 7
-    assert len(Notification.query.filter_by(notification_type=notification_type).all()) == 1
+    assert Notification.query.count() == 7
+    assert Notification.query.filter_by(notification_type=notification_type).count() == 1
     if notification_type == 'letter':
         mock_get_s3.assert_called_with(bucket_name=current_app.config['LETTERS_PDF_BUCKET_NAME'],
                                        subfolder="{}/NOTIFY.LETTER_REF.D.2.C.C".format(str(datetime.utcnow().date()))
@@ -103,17 +135,13 @@ def test_delete_notifications_for_days_of_retention(sample_service, notification
         mock_get_s3.assert_not_called()
 
 
-@pytest.mark.parametrize('notification_type', ['sms', 'email', 'letter'])
-def test_delete_notifications_inserts_notification_history(sample_service, notification_type, mocker):
-    mocker.patch("app.dao.notifications_dao.get_s3_bucket_objects")
-    create_test_data(notification_type, sample_service)
-    NotificationHistory.query.delete()
-    assert len(Notification.query.all()) == 9
-    delete_notifications_older_than_retention_by_type(notification_type)
-    assert len(Notification.query.all()) == 7
+def test_delete_notifications_inserts_notification_history(sample_service):
+    create_test_data('sms', sample_service)
+    assert Notification.query.count() == 9
+    delete_notifications_older_than_retention_by_type('sms')
+    assert Notification.query.count() == 7
 
-    history = NotificationHistory.query.all()
-    assert len(history) == 2
+    assert NotificationHistory.query.count() == 2
 
 
 def test_delete_notifications_updates_notification_history(sample_email_template, mocker):
@@ -140,59 +168,25 @@ def test_delete_notifications_updates_notification_history(sample_email_template
     assert history[0].sent_by == 'ses'
 
 
-def create_test_data(notification_type, sample_service, days_of_retention=3):
-    service_with_default_data_retention = create_service(service_name='default data retention')
-    email_template, letter_template, sms_template = _create_templates(sample_service)
-    default_email_template, default_letter_template, default_sms_template = _create_templates(
-        service_with_default_data_retention)
-    create_notification(template=email_template, status='delivered')
-    create_notification(template=sms_template, status='permanent-failure')
-    create_notification(template=letter_template, status='temporary-failure',
-                        reference='LETTER_REF', sent_at=datetime.utcnow())
-    create_notification(template=email_template, status='delivered',
-                        created_at=datetime.utcnow() - timedelta(days=4))
-    create_notification(template=sms_template, status='permanent-failure',
-                        created_at=datetime.utcnow() - timedelta(days=4))
-    create_notification(template=letter_template, status='temporary-failure',
-                        reference='LETTER_REF', sent_at=datetime.utcnow(),
-                        created_at=datetime.utcnow() - timedelta(days=4))
-    create_notification(template=default_email_template, status='delivered',
-                        created_at=datetime.utcnow() - timedelta(days=8))
-    create_notification(template=default_sms_template, status='permanent-failure',
-                        created_at=datetime.utcnow() - timedelta(days=8))
-    create_notification(template=default_letter_template, status='temporary-failure',
-                        reference='LETTER_REF', sent_at=datetime.utcnow(),
-                        created_at=datetime.utcnow() - timedelta(days=8))
-    create_service_data_retention(service_id=sample_service.id, notification_type=notification_type,
-                                  days_of_retention=days_of_retention)
-
-
-@pytest.mark.parametrize('notification_type', ['sms', 'email', 'letter'])
-def test_delete_notifications_keep_data_for_days_of_retention_is_longer(sample_service, notification_type, mocker):
-    mock_get_s3 = mocker.patch("app.dao.notifications_dao.get_s3_bucket_objects")
-    create_test_data(notification_type, sample_service, 15)
-    assert len(Notification.query.all()) == 9
-    delete_notifications_older_than_retention_by_type(notification_type)
-    assert len(Notification.query.filter_by().all()) == 8
-    assert len(Notification.query.filter_by(notification_type=notification_type).all()) == 2
-    if notification_type == 'letter':
-        assert mock_get_s3.called
-    else:
-        mock_get_s3.assert_not_called()
+def test_delete_notifications_keep_data_for_days_of_retention_is_longer(sample_service):
+    create_test_data('sms', sample_service, 15)
+    assert Notification.query.count() == 9
+    delete_notifications_older_than_retention_by_type('sms')
+    assert Notification.query.count() == 8
+    assert Notification.query.filter(Notification.notification_type == 'sms').count() == 2
 
 
 def test_delete_notifications_with_test_keys(sample_template, mocker):
     mocker.patch("app.dao.notifications_dao.get_s3_bucket_objects")
     create_notification(template=sample_template, key_type='test', created_at=datetime.utcnow() - timedelta(days=8))
-    assert len(Notification.query.all()) == 1
     delete_notifications_older_than_retention_by_type('sms')
-    assert len(Notification.query.filter_by().all()) == 0
+    assert Notification.query.count() == 0
 
 
 def test_delete_notifications_delete_notification_type_for_default_time_if_no_days_of_retention_for_type(
-        sample_service, mocker
+        sample_service
 ):
-    create_service_data_retention(service_id=sample_service.id, notification_type='sms',
+    create_service_data_retention(service=sample_service, notification_type='sms',
                                   days_of_retention=15)
     email_template, letter_template, sms_template = _create_templates(sample_service)
     create_notification(template=email_template, status='delivered')
@@ -204,10 +198,10 @@ def test_delete_notifications_delete_notification_type_for_default_time_if_no_da
                         created_at=datetime.utcnow() - timedelta(days=14))
     create_notification(template=letter_template, status='temporary-failure',
                         created_at=datetime.utcnow() - timedelta(days=14))
-    assert len(Notification.query.all()) == 6
+    assert Notification.query.count() == 6
     delete_notifications_older_than_retention_by_type('email')
-    assert len(Notification.query.filter_by().all()) == 5
-    assert len(Notification.query.filter_by(notification_type='email').all()) == 1
+    assert Notification.query.count() == 5
+    assert Notification.query.filter_by(notification_type='email').count() == 1
 
 
 def test_delete_notifications_does_try_to_delete_from_s3_when_letter_has_not_been_sent(sample_service, mocker):
@@ -220,9 +214,8 @@ def test_delete_notifications_does_try_to_delete_from_s3_when_letter_has_not_bee
     mock_get_s3.assert_not_called()
 
 
-@pytest.mark.parametrize('notification_type', ['sms'])
 @freeze_time("2016-01-10 12:00:00.000000")
-def test_should_not_delete_notification_if_history_does_not_exist(sample_service, notification_type, mocker):
+def test_should_not_delete_notification_if_history_does_not_exist(sample_service, mocker):
     mocker.patch("app.dao.notifications_dao.get_s3_bucket_objects")
     mocker.patch("app.dao.notifications_dao.insert_update_notification_history")
     with freeze_time('2016-01-01 12:00'):
@@ -231,45 +224,47 @@ def test_should_not_delete_notification_if_history_does_not_exist(sample_service
         create_notification(template=sms_template, status='delivered')
         create_notification(template=letter_template, status='temporary-failure')
     assert Notification.query.count() == 3
-    assert NotificationHistory.query.count() == 0
-    delete_notifications_older_than_retention_by_type(notification_type)
+    delete_notifications_older_than_retention_by_type('sms')
     assert Notification.query.count() == 3
     assert NotificationHistory.query.count() == 0
 
 
-def test_delete_notifications_calls_subquery(
-        notify_db_session, mocker
-):
-    service = create_service()
-    sms_template = create_template(service=service)
-    create_notification(template=sms_template, created_at=datetime.now() - timedelta(days=8))
-    create_notification(template=sms_template, created_at=datetime.now() - timedelta(days=8))
-    create_notification(template=sms_template, created_at=datetime.now() - timedelta(days=8))
+def test_delete_notifications_calls_subquery_multiple_times(sample_template):
+    create_notification(template=sample_template, created_at=datetime.now() - timedelta(days=8))
+    create_notification(template=sample_template, created_at=datetime.now() - timedelta(days=8))
+    create_notification(template=sample_template, created_at=datetime.now() - timedelta(days=8))
 
     assert Notification.query.count() == 3
     delete_notifications_older_than_retention_by_type('sms', qry_limit=1)
     assert Notification.query.count() == 0
 
 
-@pytest.mark.parametrize('notification_type', ['sms'])
-def test_insert_update_notification_history(sample_service, notification_type):
-    template = create_template(sample_service, template_type=notification_type)
+def test_delete_notifications_returns_sum_correctly(sample_template):
+    create_notification(template=sample_template, created_at=datetime.now() - timedelta(days=8))
+    create_notification(template=sample_template, created_at=datetime.now() - timedelta(days=8))
+
+    s2 = create_service(service_name='s2')
+    t2 = create_template(s2, template_type='sms')
+    create_notification(template=t2, created_at=datetime.now() - timedelta(days=8))
+    create_notification(template=t2, created_at=datetime.now() - timedelta(days=8))
+
+    ret = delete_notifications_older_than_retention_by_type('sms', qry_limit=1)
+    assert ret == 4
+
+
+def test_insert_update_notification_history(sample_service):
+    template = create_template(sample_service, template_type='sms')
     notification_1 = create_notification(template=template, created_at=datetime.utcnow() - timedelta(days=3))
     notification_2 = create_notification(template=template, created_at=datetime.utcnow() - timedelta(days=8))
     notification_3 = create_notification(template=template, created_at=datetime.utcnow() - timedelta(days=9))
-    other_types = ['sms', 'email', 'letter']
-    other_types.remove(notification_type)
+    other_types = ['email', 'letter']
     for template_type in other_types:
         t = create_template(service=sample_service, template_type=template_type)
         create_notification(template=t, created_at=datetime.utcnow() - timedelta(days=3))
         create_notification(template=t, created_at=datetime.utcnow() - timedelta(days=8))
 
-    NotificationHistory.query.delete()
-    history = NotificationHistory.query.all()
-    assert len(history) == 0
-
     insert_update_notification_history(
-        notification_type=notification_type, date_to_delete_from=datetime.utcnow() - timedelta(days=7),
+        notification_type='sms', date_to_delete_from=datetime.utcnow() - timedelta(days=7),
         service_id=sample_service.id)
     history = NotificationHistory.query.all()
     assert len(history) == 2
@@ -280,23 +275,16 @@ def test_insert_update_notification_history(sample_service, notification_type):
     assert notification_3.id in history_ids
 
 
-@pytest.mark.parametrize('notification_type', ['sms', 'email', 'letter'])
-def test_insert_update_notification_history_only_insert_update_given_service(sample_service, notification_type):
+def test_insert_update_notification_history_only_insert_update_given_service(sample_service):
     other_service = create_service(service_name='another service')
-    other_template = create_template(service=other_service, template_type=notification_type)
-    template = create_template(service=sample_service, template_type=notification_type)
+    other_template = create_template(service=other_service)
+    template = create_template(service=sample_service)
     notification_1 = create_notification(template=template, created_at=datetime.utcnow() - timedelta(days=3))
     notification_2 = create_notification(template=template, created_at=datetime.utcnow() - timedelta(days=8))
     notification_3 = create_notification(template=other_template, created_at=datetime.utcnow() - timedelta(days=3))
     notification_4 = create_notification(template=other_template, created_at=datetime.utcnow() - timedelta(days=8))
 
-    NotificationHistory.query.delete()
-    history = NotificationHistory.query.all()
-
-    assert len(history) == 0
-
-    insert_update_notification_history(
-        notification_type, datetime.utcnow() - timedelta(days=7), sample_service.id)
+    insert_update_notification_history('sms', datetime.utcnow() - timedelta(days=7), sample_service.id)
     history = NotificationHistory.query.all()
     assert len(history) == 1
 
@@ -314,13 +302,5 @@ def test_insert_update_notification_history_updates_history_with_new_status(samp
     insert_update_notification_history(
         'sms', datetime.utcnow() - timedelta(days=7), sample_template.service_id)
     history = NotificationHistory.query.get(notification_2.id)
-    assert history.id == notification_2.id
     assert history.status == 'delivered'
     assert not NotificationHistory.query.get(notification_1.id)
-
-
-def _create_templates(sample_service):
-    email_template = create_template(service=sample_service, template_type='email')
-    sms_template = create_template(service=sample_service)
-    letter_template = create_template(service=sample_service, template_type='letter')
-    return email_template, letter_template, sms_template

tests/app/dao/test_inbound_sms_dao.py

@@ -97,10 +97,10 @@ def test_should_delete_inbound_sms_according_to_data_retention(notify_db_session
 
     services = [short_retention_service, no_retention_service, long_retention_service]
 
-    create_service_data_retention(long_retention_service.id, notification_type='sms', days_of_retention=30)
-    create_service_data_retention(short_retention_service.id, notification_type='sms', days_of_retention=3)
+    create_service_data_retention(long_retention_service, notification_type='sms', days_of_retention=30)
+    create_service_data_retention(short_retention_service, notification_type='sms', days_of_retention=3)
     # email retention doesn't affect anything
-    create_service_data_retention(short_retention_service.id, notification_type='email', days_of_retention=4)
+    create_service_data_retention(short_retention_service, notification_type='email', days_of_retention=4)
 
     dates = [
         datetime(2017, 6, 4, 23, 00),  # just before three days

tests/app/dao/test_service_data_retention_dao.py

@@ -141,8 +141,8 @@ def test_update_service_data_retention_does_not_update_row_if_data_retention_is_
                           ('letter', 'sms')]
                          )
 def test_fetch_service_data_retention_by_notification_type(sample_service, notification_type, alternate):
-    data_retention = create_service_data_retention(service_id=sample_service.id, notification_type=notification_type)
-    create_service_data_retention(service_id=sample_service.id, notification_type=alternate)
+    data_retention = create_service_data_retention(service=sample_service, notification_type=notification_type)
+    create_service_data_retention(service=sample_service, notification_type=alternate)
     result = fetch_service_data_retention_by_notification_type(sample_service.id, notification_type)
     assert result == data_retention
 

tests/app/db.py

@@ -806,12 +806,12 @@ def ses_notification_callback():
 
 
 def create_service_data_retention(
-        service_id,
+        service,
         notification_type='sms',
         days_of_retention=3
 ):
     data_retention = insert_service_data_retention(
-        service_id=service_id,
+        service_id=service.id,
         notification_type=notification_type,
         days_of_retention=days_of_retention
     )

tests/app/inbound_sms/test_rest.py

@@ -113,7 +113,7 @@ def test_post_to_get_inbound_sms_for_service_respects_data_retention(
     too_old_date,
     returned_date
 ):
-    create_service_data_retention(sample_service.id, 'sms', days_of_retention)
+    create_service_data_retention(sample_service, 'sms', days_of_retention)
     create_inbound_sms(sample_service, created_at=too_old_date)
     returned_inbound = create_inbound_sms(sample_service, created_at=returned_date)
 
@@ -217,7 +217,7 @@ def test_get_most_recent_inbound_sms_for_service_respects_data_retention(
     admin_request,
     sample_service
 ):
-    create_service_data_retention(sample_service.id, 'sms', 5)
+    create_service_data_retention(sample_service, 'sms', 5)
     for i in range(10):
         created = datetime.utcnow() - timedelta(days=i)
         create_inbound_sms(sample_service, user_number='44770090000{}'.format(i), created_at=created)
@@ -240,7 +240,7 @@ def test_get_most_recent_inbound_sms_for_service_respects_data_retention_if_olde
     admin_request,
     sample_service
 ):
-    create_service_data_retention(sample_service.id, 'sms', 14)
+    create_service_data_retention(sample_service, 'sms', 14)
     create_inbound_sms(sample_service, created_at=datetime(2017, 4, 1, 12, 0))
 
     response = admin_request.get('inbound_sms.get_most_recent_inbound_sms_for_service', service_id=sample_service.id)
@@ -254,7 +254,7 @@ def test_get_inbound_sms_for_service_respects_data_retention(
     admin_request,
     sample_service
 ):
-    create_service_data_retention(sample_service.id, 'sms', 5)
+    create_service_data_retention(sample_service, 'sms', 5)
     for i in range(10):
         created = datetime.utcnow() - timedelta(days=i)
         create_inbound_sms(sample_service, user_number='44770090000{}'.format(i), created_at=created)

tests/app/service/test_service_data_retention_rest.py

@@ -7,10 +7,10 @@ from tests.app.db import create_service_data_retention
 
 
 def test_get_service_data_retention(client, sample_service):
-    sms_data_retention = create_service_data_retention(service_id=sample_service.id)
-    email_data_retention = create_service_data_retention(service_id=sample_service.id, notification_type='email',
+    sms_data_retention = create_service_data_retention(service=sample_service)
+    email_data_retention = create_service_data_retention(service=sample_service, notification_type='email',
                                                          days_of_retention=10)
-    letter_data_retention = create_service_data_retention(service_id=sample_service.id, notification_type='letter',
+    letter_data_retention = create_service_data_retention(service=sample_service, notification_type='letter',
                                                           days_of_retention=30)
 
     response = client.get(
@@ -36,7 +36,7 @@ def test_get_service_data_retention_returns_empty_list(client, sample_service):
 
 
 def test_get_data_retention_for_service_notification_type(client, sample_service):
-    data_retention = create_service_data_retention(service_id=sample_service.id)
+    data_retention = create_service_data_retention(service=sample_service)
     response = client.get('/service/{}/data-retention/notification-type/{}'.format(sample_service.id, 'sms'),
                           headers=[('Content-Type', 'application/json'), create_authorization_header()],
                           )
@@ -45,10 +45,10 @@ def test_get_data_retention_for_service_notification_type(client, sample_service
 
 
 def test_get_service_data_retention_by_id(client, sample_service):
-    sms_data_retention = create_service_data_retention(service_id=sample_service.id)
-    create_service_data_retention(service_id=sample_service.id, notification_type='email',
+    sms_data_retention = create_service_data_retention(service=sample_service)
+    create_service_data_retention(service=sample_service, notification_type='email',
                                   days_of_retention=10)
-    create_service_data_retention(service_id=sample_service.id, notification_type='letter',
+    create_service_data_retention(service=sample_service, notification_type='letter',
                                   days_of_retention=30)
     response = client.get(
         '/service/{}/data-retention/{}'.format(str(sample_service.id), sms_data_retention.id),
@@ -105,7 +105,7 @@ def test_create_service_data_retention_returns_400_when_notification_type_is_inv
 def test_create_service_data_retention_returns_400_when_data_retention_for_notification_type_already_exists(
         client, sample_service
 ):
-    create_service_data_retention(service_id=sample_service.id)
+    create_service_data_retention(service=sample_service)
     data = {
         "notification_type": "sms",
         "days_of_retention": 3
@@ -123,7 +123,7 @@ def test_create_service_data_retention_returns_400_when_data_retention_for_notif
 
 
 def test_modify_service_data_retention(client, sample_service):
-    data_retention = create_service_data_retention(service_id=sample_service.id)
+    data_retention = create_service_data_retention(service=sample_service)
     data = {
         "days_of_retention": 3
     }