store api_key_id and key_type on notification

pass through from POST /notification/<type> to the celery task
also removed a couple of asserts that can fail (based on unfrozen time comparisons)

app/celery/tasks.py

@@ -38,7 +38,8 @@ from app.dao.templates_dao import dao_get_template_by_id
 from app.models import (
     Notification,
     EMAIL_TYPE,
-    SMS_TYPE
+    SMS_TYPE,
+    KEY_TYPE_NORMAL
 )
 
 
@@ -133,7 +134,13 @@ def remove_job(job_id):
 
 
 @notify_celery.task(bind=True, name="send-sms", max_retries=5, default_retry_delay=5)
-def send_sms(self, service_id, notification_id, encrypted_notification, created_at):
+def send_sms(self,
+             service_id,
+             notification_id,
+             encrypted_notification,
+             created_at,
+             api_key_id=None,
+             key_type=KEY_TYPE_NORMAL):
     task_start = monotonic()
     notification = encryption.decrypt(encrypted_notification)
     service = dao_fetch_service_by_id(service_id)
@@ -158,7 +165,9 @@ def send_sms(self, service_id, notification_id, encrypted_notification, created_
             status='created',
             created_at=datetime.strptime(created_at, DATETIME_FORMAT),
             personalisation=notification.get('personalisation'),
-            notification_type=SMS_TYPE
+            notification_type=SMS_TYPE,
+            api_key_id=api_key_id,
+            key_type=key_type
         )
         dao_create_notification(notification_db_object, SMS_TYPE)
 
@@ -176,7 +185,13 @@ def send_sms(self, service_id, notification_id, encrypted_notification, created_
 
 
 @notify_celery.task(name="send-email")
-def send_email(service_id, notification_id, encrypted_notification, created_at, reply_to_addresses=None):
+def send_email(service_id,
+               notification_id,
+               encrypted_notification,
+               created_at,
+               reply_to_addresses=None,
+               api_key_id=None,
+               key_type=KEY_TYPE_NORMAL):
     task_start = monotonic()
     notification = encryption.decrypt(encrypted_notification)
     service = dao_fetch_service_by_id(service_id)
@@ -204,7 +219,9 @@ def send_email(service_id, notification_id, encrypted_notification, created_at,
             sent_at=sent_at,
             sent_by=provider.get_name(),
             personalisation=notification.get('personalisation'),
-            notification_type=EMAIL_TYPE
+            notification_type=EMAIL_TYPE,
+            api_key_id=api_key_id,
+            key_type=key_type
         )
 
         dao_create_notification(notification_db_object, EMAIL_TYPE)

app/notifications/rest.py

@@ -267,19 +267,33 @@ def send_notification(notification_type):
     notification_id = create_uuid()
     notification.update({"template_version": template.version})
     if notification_type == SMS_TYPE:
-        send_sms.apply_async((
+        send_sms.apply_async(
-            service_id,
+            (
-            notification_id,
+                service_id,
-            encryption.encrypt(notification),
+                notification_id,
-            datetime.utcnow().strftime(DATETIME_FORMAT)
+                encryption.encrypt(notification),
-        ), queue='sms')
+                datetime.utcnow().strftime(DATETIME_FORMAT)
+            ),
+            kwargs={
+                'api_key_id': str(api_user.id),
+                'key_type': api_user.key_type
+            },
+            queue='sms'
+        )
     else:
-        send_email.apply_async((
+        send_email.apply_async(
-            service_id,
+            (
-            notification_id,
+                service_id,
-            encryption.encrypt(notification),
+                notification_id,
-            datetime.utcnow().strftime(DATETIME_FORMAT)
+                encryption.encrypt(notification),
-        ), queue='email')
+                datetime.utcnow().strftime(DATETIME_FORMAT)
+            ),
+            kwargs={
+                'api_key_id': str(api_user.id),
+                'key_type': api_user.key_type
+            },
+            queue='email'
+        )
 
     statsd_client.incr('notifications.api.{}'.format(notification_type))
     return jsonify(

tests/app/celery/test_tasks.py

@@ -1,5 +1,5 @@
 import uuid
-from datetime import datetime
+from datetime import datetime, timedelta
 
 import pytest
 from freezegun import freeze_time
@@ -21,7 +21,7 @@ from app.celery.tasks import (
 from app.clients.email.aws_ses import AwsSesClientException
 from app.dao import jobs_dao, provider_details_dao
 from app.dao.provider_statistics_dao import get_provider_statistics
-from app.models import Notification
+from app.models import Notification, KEY_TYPE_TEAM
 from tests.app import load_example_csv
 from tests.app.conftest import (
     sample_service,
@@ -462,7 +462,7 @@ def test_should_not_send_email_if_restricted_service_and_invalid_email_address(n
         Notification.query.filter_by(id=notification_id).one()
 
 
-def test_should_send_sms_template_to_and_persist_with_job_id(sample_job, mocker):
+def test_should_send_sms_template_to_and_persist_with_job_id(sample_job, sample_api_key, mocker):
     notification = _notification_json(
         sample_job.template,
         to="+447234123123",
@@ -475,7 +475,9 @@ def test_should_send_sms_template_to_and_persist_with_job_id(sample_job, mocker)
         sample_job.service.id,
         notification_id,
         encryption.encrypt(notification),
-        datetime.utcnow().strftime(DATETIME_FORMAT)
+        datetime.utcnow().strftime(DATETIME_FORMAT),
+        api_key_id=str(sample_api_key.id),
+        key_type=KEY_TYPE_TEAM
     )
     provider_tasks.send_sms_to_provider.apply_async.assert_called_once_with(
         (sample_job.service.id,
@@ -492,9 +494,11 @@ def test_should_send_sms_template_to_and_persist_with_job_id(sample_job, mocker)
     assert persisted_notification.created_at <= datetime.utcnow()
     assert not persisted_notification.sent_by
     assert persisted_notification.job_row_number == 2
+    assert persisted_notification.api_key_id == sample_api_key.id
+    assert persisted_notification.key_type == KEY_TYPE_TEAM
 
 
-def test_should_use_email_template_and_persist(sample_email_template_with_placeholders, mocker):
+def test_should_use_email_template_and_persist(sample_email_template_with_placeholders, sample_api_key, mocker):
     notification = _notification_json(
         sample_email_template_with_placeholders,
         "my_email@my_email.com",
@@ -508,20 +512,18 @@ def test_should_use_email_template_and_persist(sample_email_template_with_placeh
 
     notification_id = uuid.uuid4()
 
-    freezer = freeze_time("2016-01-01 11:09:00.00000")
+    with freeze_time("2016-01-01 11:09:00.00000"):
-    freezer.start()
+        now = datetime.utcnow()
-    now = datetime.utcnow()
-    freezer.stop()
 
-    freezer = freeze_time("2016-01-01 11:10:00.00000")
+    with freeze_time("2016-01-01 11:10:00.00000"):
-    freezer.start()
+        send_email(
-    send_email(
+            sample_email_template_with_placeholders.service_id,
-        sample_email_template_with_placeholders.service_id,
+            notification_id,
-        notification_id,
+            encryption.encrypt(notification),
-        encryption.encrypt(notification),
+            now.strftime(DATETIME_FORMAT),
-        now.strftime(DATETIME_FORMAT)
+            api_key_id=str(sample_api_key.id),
-    )
+            key_type=KEY_TYPE_TEAM
-    freezer.stop()
+        )
 
     aws_ses_client.send_email.assert_called_once_with(
         '"Sample service" <sample.service@test.notify.com>',
@@ -552,6 +554,8 @@ def test_should_use_email_template_and_persist(sample_email_template_with_placeh
     assert persisted_notification.job_row_number == 1
     assert persisted_notification.personalisation == {'name': 'Jo'}
     assert persisted_notification._personalisation == encryption.encrypt({"name": "Jo"})
+    assert persisted_notification.api_key_id == sample_api_key.id
+    assert persisted_notification.key_type == KEY_TYPE_TEAM
 
 
 def test_send_email_should_use_template_version_from_job_not_latest(sample_email_template, mocker):
@@ -589,8 +593,6 @@ def test_send_email_should_use_template_version_from_job_not_latest(sample_email
     assert persisted_notification.to == 'my_email@my_email.com'
     assert persisted_notification.template_id == sample_email_template.id
     assert persisted_notification.template_version == version_on_notification
-    assert persisted_notification.created_at == now
-    assert persisted_notification.sent_at > now
     assert persisted_notification.status == 'sending'
     assert persisted_notification.sent_by == 'ses'
 
@@ -621,8 +623,6 @@ def test_should_use_email_template_subject_placeholders(sample_email_template_wi
     assert persisted_notification.id == notification_id
     assert persisted_notification.to == 'my_email@my_email.com'
     assert persisted_notification.template_id == sample_email_template_with_placeholders.id
-    assert persisted_notification.created_at == now
-    assert persisted_notification.sent_at > now
     assert persisted_notification.status == 'sending'
     assert persisted_notification.sent_by == 'ses'
     assert persisted_notification.personalisation == {"name": "Jo"}
@@ -635,13 +635,16 @@ def test_should_use_email_template_and_persist_without_personalisation(sample_em
     mocker.patch('app.aws_ses_client.get_name', return_value='ses')
 
     notification_id = uuid.uuid4()
-    now = datetime.utcnow()
+
-    send_email(
+    with freeze_time("2016-01-01 11:09:00.00000") as freezer:
-        sample_email_template.service_id,
+        now = datetime.utcnow()
-        notification_id,
+        freezer.tick(delta=timedelta(seconds=60))
-        encryption.encrypt(notification),
+        send_email(
-        now.strftime(DATETIME_FORMAT)
+            sample_email_template.service_id,
-    )
+            notification_id,
+            encryption.encrypt(notification),
+            now.strftime(DATETIME_FORMAT)
+        )
     aws_ses_client.send_email.assert_called_once_with(
         '"Sample service" <sample.service@test.notify.com>',
         "my_email@my_email.com",
@@ -655,7 +658,7 @@ def test_should_use_email_template_and_persist_without_personalisation(sample_em
     assert persisted_notification.to == 'my_email@my_email.com'
     assert persisted_notification.template_id == sample_email_template.id
     assert persisted_notification.created_at == now
-    assert persisted_notification.sent_at > now
+    assert persisted_notification.sent_at == now + timedelta(seconds=60)
     assert persisted_notification.status == 'sending'
     assert persisted_notification.sent_by == 'ses'
     assert not persisted_notification.personalisation
@@ -802,8 +805,6 @@ def test_should_call_send_email_response_task_if_research_mode(
     assert persisted_notification.to == 'john@smith.com'
     assert persisted_notification.template_id == sample_email_template.id
     assert persisted_notification.status == 'sending'
-    assert persisted_notification.sent_at > now
-    assert persisted_notification.created_at == now
     assert persisted_notification.sent_by == 'ses'
     assert persisted_notification.reference == str(reference)
 

tests/app/notifications/rest/test_send_notification.py

@@ -5,12 +5,14 @@ import string
 from unittest.mock import ANY
 from flask import (json, current_app)
 from freezegun import freeze_time
+from notifications_python_client.authentication import create_jwt_token
 
 import app
 from app import encryption
-from app.models import KEY_TYPE_TEAM
+from app.models import ApiKey, KEY_TYPE_TEAM
 from app.dao.templates_dao import dao_get_all_templates_for_service, dao_update_template
 from app.dao.services_dao import dao_update_service
+from app.dao.api_key_dao import save_model_api_key
 from tests import create_authorization_header
 from tests.app.conftest import (
     sample_notification as create_sample_notification,
@@ -118,6 +120,7 @@ def test_send_notification_with_placeholders_replaced(notify_api, sample_email_t
                  notification_id,
                  ANY,
                  "2016-01-01T11:09:00.061258"),
+                kwargs=ANY,
                 queue="email"
             )
             assert response.status_code == 201
@@ -358,6 +361,7 @@ def test_should_allow_valid_sms_notification(notify_api, sample_template, mocker
                  notification_id,
                  "something_encrypted",
                  "2016-01-01T11:09:00.061258"),
+                kwargs=ANY,
                 queue="sms"
             )
             assert response.status_code == 201
@@ -523,6 +527,7 @@ def test_should_allow_valid_email_notification(notify_api, sample_email_template
                  notification_id,
                  "something_encrypted",
                  "2016-01-01T11:09:00.061258"),
+                kwargs=ANY,
                 queue="email"
             )
 
@@ -769,15 +774,25 @@ def test_should_send_email_if_team_api_key_and_a_service_user(notify_api, sample
             'to': sample_email_template.service.created_by.email_address,
             'template': sample_email_template.id
         }
-
+        api_key = ApiKey(service=sample_email_template.service,
-        auth_header = create_authorization_header(service_id=sample_email_template.service_id, key_type=KEY_TYPE_TEAM)
+                         name='team_key',
+                         created_by=sample_email_template.created_by,
+                         key_type=KEY_TYPE_TEAM)
+        save_model_api_key(api_key)
+        auth_header = create_jwt_token(secret=api_key.unsigned_secret, client_id=str(api_key.service_id))
 
         response = client.post(
             path='/notifications/email',
             data=json.dumps(data),
-            headers=[('Content-Type', 'application/json'), auth_header])
+            headers=[('Content-Type', 'application/json'), ('Authorization', 'Bearer {}'.format(auth_header))])
 
-        assert app.celery.tasks.send_email.apply_async.called
+        app.celery.tasks.send_email.apply_async.assert_called_once_with(
+            ANY,
+            kwargs={
+                'api_key_id': str(api_key.id),
+                'key_type': api_key.key_type
+            },
+            queue='email')
         assert response.status_code == 201
 
 
@@ -789,13 +804,23 @@ def test_should_send_sms_if_team_api_key_and_a_service_user(notify_api, sample_t
             'to': sample_template.service.created_by.mobile_number,
             'template': sample_template.id
         }
-
+        api_key = ApiKey(service=sample_template.service,
-        auth_header = create_authorization_header(service_id=sample_template.service_id, key_type=KEY_TYPE_TEAM)
+                         name='team_key',
+                         created_by=sample_template.created_by,
+                         key_type=KEY_TYPE_TEAM)
+        save_model_api_key(api_key)
+        auth_header = create_jwt_token(secret=api_key.unsigned_secret, client_id=str(api_key.service_id))
 
         response = client.post(
             path='/notifications/sms',
             data=json.dumps(data),
-            headers=[('Content-Type', 'application/json'), auth_header])
+            headers=[('Content-Type', 'application/json'), ('Authorization', 'Bearer {}'.format(auth_header))])
 
-        assert app.celery.tasks.send_sms.apply_async.called
+        app.celery.tasks.send_sms.apply_async.assert_called_once_with(
+            ANY,
+            kwargs={
+                'api_key_id': str(api_key.id),
+                'key_type': api_key.key_type
+            },
+            queue='sms')
         assert response.status_code == 201