fix more

tests/app/celery/test_tasks.py

@@ -531,7 +531,12 @@ def test_should_not_save_sms_if_restricted_service_and_invalid_number(
         encryption.encrypt(notification),
     )
     assert provider_tasks.deliver_sms.apply_async.called is False
-    assert Notification.query.count() == 0
+    assert _get_notification_query_count() == 0
+
+
+def _get_notification_query_count():
+    stmt = select(func.count()).select_from(Notification)
+    return db.session.execute(stmt).scalar() or 0
 
 
 def test_should_not_save_email_if_restricted_service_and_invalid_email_address(
@@ -553,7 +558,7 @@ def test_should_not_save_email_if_restricted_service_and_invalid_email_address(
         encryption.encrypt(notification),
     )
 
-    assert Notification.query.count() == 0
+    assert _get_notification_query_count() == 0
 
 
 def test_should_save_sms_template_to_and_persist_with_job_id(sample_job, mocker):
@@ -593,7 +598,7 @@ def test_should_save_sms_template_to_and_persist_with_job_id(sample_job, mocker)
 def test_should_not_save_sms_if_team_key_and_recipient_not_in_team(
     notify_db_session, mocker
 ):
-    assert Notification.query.count() == 0
+    assert _get_notification_query_count() == 0
     user = create_user(mobile_number="2028675309")
     service = create_service(user=user, restricted=True)
     template = create_template(service=service)
@@ -611,7 +616,7 @@ def test_should_not_save_sms_if_team_key_and_recipient_not_in_team(
         encryption.encrypt(notification),
     )
     assert provider_tasks.deliver_sms.apply_async.called is False
-    assert Notification.query.count() == 0
+    assert _get_notification_query_count() == 0
 
 
 def test_should_use_email_template_and_persist(
@@ -836,7 +841,7 @@ def test_save_sms_should_go_to_retry_queue_if_database_errors(sample_template, m
     assert provider_tasks.deliver_sms.apply_async.called is False
     tasks.save_sms.retry.assert_called_with(exc=expected_exception, queue="retry-tasks")
 
-    assert Notification.query.count() == 0
+    assert _get_notification_query_count() == 0
 
 
 def test_save_email_should_go_to_retry_queue_if_database_errors(
@@ -866,7 +871,7 @@ def test_save_email_should_go_to_retry_queue_if_database_errors(
         exc=expected_exception, queue="retry-tasks"
     )
 
-    assert Notification.query.count() == 0
+    assert _get_notification_query_count() == 0
 
 
 def test_save_email_does_not_send_duplicate_and_does_not_put_in_retry_queue(
@@ -888,7 +893,7 @@ def test_save_email_does_not_send_duplicate_and_does_not_put_in_retry_queue(
         notification_id,
         encryption.encrypt(json),
     )
-    assert Notification.query.count() == 1
+    assert _get_notification_query_count() == 1
     assert not deliver_email.called
     assert not retry.called
 
@@ -912,7 +917,7 @@ def test_save_sms_does_not_send_duplicate_and_does_not_put_in_retry_queue(
         notification_id,
         encryption.encrypt(json),
     )
-    assert Notification.query.count() == 1
+    assert _get_notification_query_count() == 1
     assert not deliver_sms.called
     assert not retry.called
 

tests/app/dao/test_invited_user_dao.py

@@ -19,8 +19,13 @@ from app.utils import utc_now
 from tests.app.db import create_invited_user
 
 
+def _get_invited_user_count():
+    stmt = select(func.count()).select_from(InvitedUser)
+    return db.session.execute(stmt).scalar() or 0
+
+
 def test_create_invited_user(notify_db_session, sample_service):
-    assert InvitedUser.query.count() == 0
+    assert _get_invited_user_count() == 0
     email_address = "invited_user@service.gov.uk"
     invite_from = sample_service.users[0]
 
@@ -35,7 +40,7 @@ def test_create_invited_user(notify_db_session, sample_service):
     invited_user = InvitedUser(**data)
     save_invited_user(invited_user)
 
-    assert InvitedUser.query.count() == 1
+    assert _get_invited_user_count() == 1
     assert invited_user.email_address == email_address
     assert invited_user.from_user == invite_from
     permissions = invited_user.get_permissions()
@@ -48,7 +53,7 @@ def test_create_invited_user(notify_db_session, sample_service):
 def test_create_invited_user_sets_default_folder_permissions_of_empty_list(
     sample_service,
 ):
-    assert InvitedUser.query.count() == 0
+    assert _get_invited_user_count() == 0
     invite_from = sample_service.users[0]
 
     data = {
@@ -61,7 +66,7 @@ def test_create_invited_user_sets_default_folder_permissions_of_empty_list(
     invited_user = InvitedUser(**data)
     save_invited_user(invited_user)
 
-    assert InvitedUser.query.count() == 1
+    assert _get_invited_user_count() == 1
     assert invited_user.folder_permissions == []
 
 
@@ -109,12 +114,12 @@ def test_get_invited_users_for_service_that_has_no_invites(
 def test_save_invited_user_sets_status_to_cancelled(
     notify_db_session, sample_invited_user
 ):
-    assert InvitedUser.query.count() == 1
+    assert _get_invited_user_count() == 1
     saved = InvitedUser.query.get(sample_invited_user.id)
     assert saved.status == InvitedUserStatus.PENDING
     saved.status = InvitedUserStatus.CANCELLED
     save_invited_user(saved)
-    assert InvitedUser.query.count() == 1
+    assert _get_invited_user_count() == 1
     cancelled_invited_user = InvitedUser.query.get(sample_invited_user.id)
     assert cancelled_invited_user.status == InvitedUserStatus.CANCELLED
 

tests/app/service/test_api_key_endpoints.py

@@ -1,7 +1,9 @@
 import json
 
 from flask import url_for
+from sqlalchemy import func, select
 
+from app import db
 from app.dao.api_key_dao import expire_api_key
 from app.enums import KeyType
 from app.models import ApiKey
@@ -60,10 +62,15 @@ def test_create_api_key_without_key_type_rejects(client, sample_service):
     assert json_resp["message"] == {"key_type": ["Missing data for required field."]}
 
 
+def _get_api_key_count():
+    stmt = select(func.count()).select_from(ApiKey)
+    return db.session.execute(stmt).scalar() or 0
+
+
 def test_revoke_should_expire_api_key_for_service(notify_api, sample_api_key):
     with notify_api.test_request_context():
         with notify_api.test_client() as client:
-            assert ApiKey.query.count() == 1
+            assert _get_api_key_count() == 1
             auth_header = create_admin_authorization_header()
             response = client.post(
                 url_for(
@@ -83,7 +90,7 @@ def test_api_key_should_create_multiple_new_api_key_for_service(
 ):
     with notify_api.test_request_context():
         with notify_api.test_client() as client:
-            assert ApiKey.query.count() == 0
+            assert _get_api_key_count() == 0
             data = {
                 "name": "some secret name",
                 "created_by": str(sample_service.created_by.id),
@@ -96,7 +103,7 @@ def test_api_key_should_create_multiple_new_api_key_for_service(
                 headers=[("Content-Type", "application/json"), auth_header],
             )
             assert response.status_code == 201
-            assert ApiKey.query.count() == 1
+            assert _get_api_key_count() == 1
 
             data["name"] = "another secret name"
             auth_header = create_admin_authorization_header()
@@ -109,7 +116,7 @@ def test_api_key_should_create_multiple_new_api_key_for_service(
             assert json.loads(response.get_data(as_text=True)) != json.loads(
                 response2.get_data(as_text=True)
             )
-            assert ApiKey.query.count() == 2
+            assert _get_api_key_count() == 2
 
 
 def test_get_api_keys_should_return_all_keys_for_service(notify_api, sample_api_key):
@@ -130,7 +137,7 @@ def test_get_api_keys_should_return_all_keys_for_service(notify_api, sample_api_
                 service_id=one_to_expire.service_id, api_key_id=one_to_expire.id
             )
 
-            assert ApiKey.query.count() == 4
+            assert _get_api_key_count() == 4
 
             auth_header = create_admin_authorization_header()
             response = client.get(

tests/app/user/test_rest_verify.py

@@ -5,6 +5,7 @@ from datetime import datetime, timedelta
 import pytest
 from flask import current_app, url_for
 from freezegun import freeze_time
+from sqlalchemy import func, select
 
 import app.celery.tasks
 from app import db
@@ -295,7 +296,7 @@ def test_send_sms_code_returns_204_when_too_many_codes_already_created(
         )
         db.session.add(verify_code)
         db.session.commit()
-    assert VerifyCode.query.count() == 5
+    assert _get_verify_code_count() == 5
     auth_header = create_admin_authorization_header()
     resp = client.post(
         url_for(
@@ -307,7 +308,12 @@ def test_send_sms_code_returns_204_when_too_many_codes_already_created(
         headers=[("Content-Type", "application/json"), auth_header],
     )
     assert resp.status_code == 204
-    assert VerifyCode.query.count() == 5
+    assert _get_verify_code_count() == 5
+
+
+def _get_verify_code_count():
+    stmt = select(func.count()).select_from(VerifyCode)
+    return db.session.execute(stmt).scalar() or 0
 
 
 @pytest.mark.parametrize(
@@ -341,7 +347,7 @@ def test_send_new_user_email_verification(
     notify_service = email_verification_template.service
     assert resp.status_code == 204
     notification = Notification.query.first()
-    assert VerifyCode.query.count() == 0
+    assert _get_verify_code_count() == 0
     mocked.assert_called_once_with(
         ([str(notification.id)]), queue="notify-internal-tasks"
     )