fix more

tests/app/celery/test_tasks.py

@@ -539,6 +539,11 @@ def test_should_not_save_sms_if_restricted_service_and_invalid_number(
     assert _get_notification_query_count() == 0
 
 
+def _get_notification_query_all():
+    stmt = select(Notification)
+    return db.session.execute(stmt).scalars().all()
+
+
 def _get_notification_query_count():
     stmt = select(func.count()).select_from(Notification)
     return db.session.execute(stmt).scalar() or 0
@@ -1481,12 +1486,12 @@ def test_save_api_email_or_sms(mocker, sample_service, notification_type):
 
     encrypted = encryption.encrypt(data)
 
-    assert len(Notification.query.all()) == 0
+    assert len(_get_notification_query_all()) == 0
     if notification_type == NotificationType.EMAIL:
         save_api_email(encrypted_notification=encrypted)
     else:
         save_api_sms(encrypted_notification=encrypted)
-    notifications = Notification.query.all()
+    notifications = _get_notification_query_all()
     assert len(notifications) == 1
     assert str(notifications[0].id) == data["id"]
     assert notifications[0].created_at == datetime(2020, 3, 25, 14, 30)
@@ -1534,20 +1539,20 @@ def test_save_api_email_dont_retry_if_notification_already_exists(
         expected_queue = QueueNames.SEND_SMS
 
     encrypted = encryption.encrypt(data)
-    assert len(Notification.query.all()) == 0
+    assert len(_get_notification_query_all()) == 0
 
     if notification_type == NotificationType.EMAIL:
         save_api_email(encrypted_notification=encrypted)
     else:
         save_api_sms(encrypted_notification=encrypted)
-    notifications = Notification.query.all()
+    notifications = _get_notification_query_all()
     assert len(notifications) == 1
     # call the task again with the same notification
     if notification_type == NotificationType.EMAIL:
         save_api_email(encrypted_notification=encrypted)
     else:
         save_api_sms(encrypted_notification=encrypted)
-    notifications = Notification.query.all()
+    notifications = _get_notification_query_all()
     assert len(notifications) == 1
     assert str(notifications[0].id) == data["id"]
     assert notifications[0].created_at == datetime(2020, 3, 25, 14, 30)
@@ -1611,7 +1616,7 @@ def test_save_tasks_use_cached_service_and_template(
     ]
 
     # But we save 2 notifications and enqueue 2 tasks
-    assert len(Notification.query.all()) == 2
+    assert len(_get_notification_query_all()) == 2
     assert len(delivery_mock.call_args_list) == 2
 
 
@@ -1672,12 +1677,12 @@ def test_save_api_tasks_use_cache(
             }
         )
 
-    assert len(Notification.query.all()) == 0
+    assert len(_get_notification_query_all()) == 0
 
     for _ in range(3):
         task_function(encrypted_notification=create_encrypted_notification())
 
     assert service_dict_mock.call_args_list == [call(str(template.service_id))]
 
-    assert len(Notification.query.all()) == 3
+    assert len(_get_notification_query_all()) == 3
     assert len(mock_provider_task.call_args_list) == 3

tests/app/dao/notification_dao/test_notification_dao_delete_notifications.py

@@ -43,11 +43,21 @@ def test_move_notifications_does_nothing_if_notification_history_row_already_exi
     )
 
     assert _get_notification_count() == 0
-    history = NotificationHistory.query.all()
+    history = _get_notification_history_query_all()
     assert len(history) == 1
     assert history[0].status == NotificationStatus.DELIVERED
 
 
+def _get_notification_query_all():
+    stmt = select(Notification)
+    return db.session.execute(stmt).scalars().all()
+
+
+def _get_notification_history_query_all():
+    stmt = select(NotificationHistory)
+    return db.session.execute(stmt).scalars().all()
+
+
 def _get_notification_count():
     stmt = select(func.count()).select_from(Notification)
     return db.session.execute(stmt).scalar() or 0
@@ -76,8 +86,18 @@ def test_move_notifications_only_moves_notifications_older_than_provided_timesta
     )
     assert result == 1
 
-    assert Notification.query.one().id == new_notification.id
-    assert NotificationHistory.query.one().id == old_notification_id
+    assert _get_notification_query_one().id == new_notification.id
+    assert _get_notification_history_query_one().id == old_notification_id
+
+
+def _get_notification_query_one():
+    stmt = select(Notification)
+    return db.session.execute(stmt).scalars().one()
+
+
+def _get_notification_history_query_one():
+    stmt = select(NotificationHistory)
+    return db.session.execute(stmt).scalars().one()
 
 
 def test_move_notifications_keeps_calling_until_no_more_to_delete_and_then_returns_total_deleted(
@@ -123,7 +143,9 @@ def test_move_notifications_only_moves_for_given_notification_type(sample_servic
     )
     assert result == 1
     assert {x.notification_type for x in Notification.query} == {NotificationType.EMAIL}
-    assert NotificationHistory.query.one().notification_type == NotificationType.SMS
+    assert (
+        _get_notification_history_query_one().notification_type == NotificationType.SMS
+    )
 
 
 def test_move_notifications_only_moves_for_given_service(notify_db_session):
@@ -146,8 +168,8 @@ def test_move_notifications_only_moves_for_given_service(notify_db_session):
     )
     assert result == 1
 
-    assert NotificationHistory.query.one().service_id == service.id
-    assert Notification.query.one().service_id == other_service.id
+    assert _get_notification_history_query_one().service_id == service.id
+    assert _get_notification_query_one().service_id == other_service.id
 
 
 def test_move_notifications_just_deletes_test_key_notifications(sample_template):
@@ -258,8 +280,8 @@ def test_insert_notification_history_delete_notifications(sample_email_template)
         timestamp_to_delete_backwards_from=utc_now() - timedelta(days=1),
     )
     assert del_count == 8
-    notifications = Notification.query.all()
-    history_rows = NotificationHistory.query.all()
+    notifications = _get_notification_query_all()
+    history_rows = _get_notification_history_query_all()
     assert len(history_rows) == 8
     assert ids_to_move == sorted([x.id for x in history_rows])
     assert len(notifications) == 3
@@ -293,8 +315,8 @@ def test_insert_notification_history_delete_notifications_more_notifications_tha
     )
 
     assert del_count == 1
-    notifications = Notification.query.all()
-    history_rows = NotificationHistory.query.all()
+    notifications = _get_notification_query_all()
+    history_rows = _get_notification_history_query_all()
     assert len(history_rows) == 1
     assert len(notifications) == 2
 
@@ -324,8 +346,8 @@ def test_insert_notification_history_delete_notifications_only_insert_delete_for
     )
 
     assert del_count == 1
-    notifications = Notification.query.all()
-    history_rows = NotificationHistory.query.all()
+    notifications = _get_notification_query_all()
+    history_rows = _get_notification_history_query_all()
     assert len(notifications) == 1
     assert len(history_rows) == 1
     assert notifications[0].id == notification_to_stay.id
@@ -361,8 +383,8 @@ def test_insert_notification_history_delete_notifications_insert_for_key_type(
     )
 
     assert del_count == 2
-    notifications = Notification.query.all()
-    history_rows = NotificationHistory.query.all()
+    notifications = _get_notification_query_all()
+    history_rows = _get_notification_history_query_all()
     assert len(notifications) == 1
     assert with_test_key.id == notifications[0].id
     assert len(history_rows) == 2

tests/app/test_commands.py

@@ -177,8 +177,7 @@ def test_populate_organization_agreement_details_from_file(
     org_count = _get_organization_query_count()
     assert org_count == 1
 
-    org = Organization.query.one()
-    org.agreement_signed = True
+    org = _get_organization_query_one()
     notify_db_session.commit()
 
     text = (
@@ -195,11 +194,16 @@ def test_populate_organization_agreement_details_from_file(
 
     org_count = _get_organization_query_count()
     assert org_count == 1
-    org = Organization.query.one()
+    org = _get_organization_query_one()
     assert org.agreement_signed_on_behalf_of_name == "bob"
     os.remove(file_name)
 
 
+def _get_organization_query_one():
+    stmt = select(Organization)
+    return db.session.execute(stmt).scalars().one()
+
+
 def test_bulk_invite_user_to_service(
     notify_db_session, notify_api, sample_service, sample_user
 ):
@@ -344,9 +348,14 @@ def test_populate_annual_billing_with_the_previous_years_allowance(
     assert results[0].free_sms_fragment_limit == expected_allowance
 
 
+def _get_notification_query_one():
+    stmt = select(Notification)
+    return db.session.execute(stmt).scalars().one()
+
+
 def test_fix_billable_units(notify_db_session, notify_api, sample_template):
     create_notification(template=sample_template)
-    notification = Notification.query.one()
+    notification = _get_notification_query_one()
     notification.billable_units = 0
     notification.notification_type = NotificationType.SMS
     notification.status = NotificationStatus.DELIVERED
@@ -357,7 +366,7 @@ def test_fix_billable_units(notify_db_session, notify_api, sample_template):
 
     notify_api.test_cli_runner().invoke(fix_billable_units, [])
 
-    notification = Notification.query.one()
+    notification = _get_notification_query_one()
     assert notification.billable_units == 1