from datetime import UTC, datetime, timedelta from unittest.mock import ANY, patch from django.conf import settings from django.contrib.auth import get_user_model from django.core.files.base import ContentFile from django.test import TestCase from rest_framework.test import APITestCase from letters.models import LetterImage from .models import Letter User = get_user_model() class LetterModelTest(TestCase): def setUp(self): self.user = User.objects.create_user(email="test@pi-ku.app", password="password1234", full_name="Test User") def test_create_letter_draft(self): """ Test the Letter model is created with required fields and auto timestamps. """ letter = Letter.objects.create(user=self.user, type="KEPT", status="DRAFT") self.assertEqual(letter.user, self.user) self.assertEqual(letter.type, "KEPT") self.assertEqual(letter.status, "DRAFT") self.assertIsNotNone(letter.public_id) self.assertIsNone(letter.unlock_at) self.assertEqual(letter.encrypted_content, None) self.assertEqual(letter.encrypted_metadata, None) self.assertIsNone(letter.sealed_at) self.assertIsNone(letter.opened_at) self.assertIsNone(letter.burned_at) self.assertIsNotNone(letter.created_at) self.assertIsNotNone(letter.updated_at) def test_vault_requires_unlock_date_when_sealed(self): """ Test that a sealed VAULT letter cannot be created without an unlock_date """ from django.core.exceptions import ValidationError letter = Letter( user=self.user, type=Letter.Type.VAULT, status=Letter.Status.SEALED, encrypted_content="enc_content==", encrypted_metadata="enc_meta==", encrypted_dek="enc_dek==", ) with self.assertRaises(ValidationError): letter.full_clean() class LetterAPITest(APITestCase): def setUp(self): self.user = User.objects.create_user(email="api@pi-ku.app", password="password1234", full_name="API User") self.client.force_authenticate(user=self.user) self.url = "/api/letters/" def test_create_draft_letter_api(self): """ Test that the API can successfully create a basic draft letter. """ payload = { "public_id": "4281edcc-5459-4ff2-bb5e-669fb44e0757", "type": "KEPT", "encrypted_content": "enc_xyz==", "encrypted_metadata": "enc_meta==", "encrypted_dek": "enc_dek==", } response = self.client.put(self.url + payload["public_id"] + "/", payload) self.assertEqual(response.status_code, 201) self.assertEqual(Letter.objects.count(), 1) self.assertEqual(Letter.objects.get().status, "DRAFT") self.assertEqual(Letter.objects.get().type, "KEPT") self.assertEqual(Letter.objects.get().user, self.user) def test_update_draft_letter_with_public_id(self): """ Test API can successfully update an existing letter with new values. """ letter = Letter.objects.create( user=self.user, type="KEPT", status="DRAFT", public_id="4281edcc-5459-4ff2-bb5e-669fb44e0757", encrypted_content="enc_xyz==", encrypted_metadata="enc_meta==", encrypted_dek="enc_dek==", ) payload = { "public_id": letter.public_id, "type": "KEPT", "encrypted_content": "enc_abc==", "encrypted_metadata": "enc_meta==", "encrypted_dek": "enc_dek==", } response = self.client.put(self.url + letter.public_id + "/", payload) self.assertEqual(response.status_code, 200) self.assertEqual(Letter.objects.count(), 1) self.assertEqual(Letter.objects.get().status, "DRAFT") self.assertEqual(Letter.objects.get().type, "KEPT") self.assertEqual(Letter.objects.get().user, self.user) self.assertEqual(Letter.objects.get().encrypted_content, "enc_abc==") self.assertEqual(Letter.objects.get().encrypted_metadata, "enc_meta==") self.assertEqual(Letter.objects.get().encrypted_dek, "enc_dek==") def test_sealed_letters_cannot_be_updated(self): """ Test that the API returns 400 when a user tries to update an already sealed letter. """ letter = Letter.objects.create( user=self.user, type="KEPT", status="SEALED", public_id="4281edcc-5459-4ff2-bb5e-669fb44e0757", encrypted_content="enc_xyz==", encrypted_metadata="enc_meta==", encrypted_dek="enc_dek==", ) payload = { "public_id": letter.public_id, "type": "KEPT", "encrypted_content": "enc_abc==", "encrypted_metadata": "enc_meta==", "encrypted_dek": "enc_dek==", } response = self.client.put(self.url + letter.public_id + "/", payload) self.assertEqual(response.status_code, 400) self.assertEqual(response.data, {"error": "Sealed letters cannot be modified."}) def test_encrypted_dek_is_required_when_storing_encrypted_content_and_metadata(self): """ Test that encrypted_dek is required when encrypted_content and encrypted_metadata are added to the letter. """ payload = {"type": "KEPT", "encrypted_content": "enc_xyz==", "encrypted_metadata": "enc_meta=="} response = self.client.post(self.url, payload) self.assertEqual(response.status_code, 400) self.assertEqual(Letter.objects.count(), 0) self.assertEqual( response.data["non_field_errors"], ["encrypted_dek is required when encrypted_content and encrypted_metadata are present"], ) def test_create_letter_with_images_api(self): """ Test that the API can create a letter and attach encrypted images in one request. """ from django.core.files.uploadedfile import SimpleUploadedFile # Simulate local files upload image1 = SimpleUploadedFile("enc_img1.bin", b"encrypted_bytes_1", content_type="application/octet-stream") image2 = SimpleUploadedFile("enc_img2.bin", b"encrypted_bytes_2", content_type="application/octet-stream") payload = { "public_id": "4281edcc-5459-4ff2-bb5e-669fb44e0757", "type": "SENT", "status": "SEALED", "encrypted_content": "enc_content==", "encrypted_metadata": "enc_metadata==", "encrypted_dek": "enc_dek==", "image_files": [image1, image2], } response = self.client.put(self.url + payload["public_id"] + "/", payload, format="multipart") self.assertEqual(response.status_code, 201) self.assertEqual(Letter.objects.count(), 1) self.assertEqual(LetterImage.objects.count(), 2) from django.core.files.storage import default_storage self.assertTrue(default_storage.exists("encrypted-images/enc_img1.bin")) self.assertTrue(default_storage.exists("encrypted-images/enc_img2.bin")) def test_cleanup_images_when_letter_is_updated(self): """ Test that the old images are cleaned up when a letter is updated with new images. """ letter = Letter.objects.create(user=self.user, type="KEPT", status="DRAFT") LetterImage.objects.create(letter=letter, file_name="old1.bin", file=ContentFile(b"data", name="old1.bin")) LetterImage.objects.create(letter=letter, file_name="old2.bin", file=ContentFile(b"data", name="old2.bin")) response = self.client.put( f"/api/letters/{letter.public_id}/", data={ "encrypted_content": "new_enc==", "encrypted_metadata": "new_meta==", "encrypted_dek": "new_dek==", "image_files": [ContentFile(b"data", name="new.bin")], }, format="multipart", ) from django.core.files.storage import default_storage # Verify that the old files are cleared from storage directory as well self.assertTrue(LetterImage.objects.filter(file_name="new.bin").exists()) self.assertEqual(LetterImage.objects.count(), 1) self.assertFalse(default_storage.exists("encrypted-images/old1.bin")) self.assertFalse(default_storage.exists("encrypted-images/old2.bin")) self.assertEqual(response.status_code, 200) def test_vault_letters_does_not_return_letter_content_before_the_unlock_date(self): """ Test that the vault letters does not return letter content (images and encrypted_content) before the unlock date. """ from datetime import datetime, timedelta letter = Letter.objects.create( user=self.user, type="VAULT", status="SEALED", public_id="4281edcc-5459-4ff2-bb5e-669fb44e0757", encrypted_content="enc_content==", encrypted_metadata="enc_meta==", encrypted_dek="enc_dek==", unlock_at=datetime.now(UTC), ) from freezegun import freeze_time past_datetime = datetime.now(UTC) - timedelta(seconds=1) future_datetime = datetime.now(UTC) + timedelta(seconds=1) with freeze_time(past_datetime): response = self.client.get(f"/api/letters/{letter.public_id}/") self.assertEqual(response.status_code, 200) self.assertEqual(response.data["encrypted_content"], None) self.assertEqual(response.data["encrypted_metadata"], "enc_meta==") self.assertEqual(response.data["encrypted_dek"], None) with freeze_time(future_datetime): response = self.client.get(f"/api/letters/{letter.public_id}/") self.assertEqual(response.status_code, 200) self.assertEqual(response.data["encrypted_content"], "enc_content==") self.assertEqual(response.data["encrypted_metadata"], "enc_meta==") self.assertEqual(response.data["encrypted_dek"], "enc_dek==") class LetterImageModelTest(TestCase): def setUp(self): self.user = User.objects.create_user(email="img_test@pi-ku.app", password="password1234") self.letter = Letter.objects.create(user=self.user, type="KEPT", status="DRAFT") def test_create_letter_image(self): """ Test that images can be associated with a letter (many to 1). """ image_content = ContentFile(b"fake-encrypted-data", name="test_image.bin") letter_image = LetterImage.objects.create( letter=self.letter, file_name="encrypted_image.enc", file=image_content ) self.assertEqual(letter_image.letter, self.letter) self.assertTrue(letter_image.file.name.startswith("encrypted-images/")) self.assertIsNotNone(letter_image.public_id) def test_letter_cascade_deletes_images(self): """ TTest that when a letter is deleted, its encrypted images are also removed. """ LetterImage.objects.create( letter=self.letter, file_name="will_be_deleted.jpg", file=ContentFile(b"data", name="del.bin") ) self.assertEqual(LetterImage.objects.count(), 1) self.letter.delete() self.assertEqual(LetterImage.objects.count(), 0) class LetterTaskTest(TestCase): def setUp(self): self.user = User.objects.create_user(email="task@pi-ku.app", password="password1234") def test_get_vault_letters_to_be_notified(self): """ Test that the task can successfully retrieve the letters whose unlock date is passed and haven't been notified. """ from letters.tasks import get_vault_letters_to_notify Letter.objects.create( user=self.user, type="VAULT", status="SEALED", unlock_at=datetime.now(UTC) + timedelta(seconds=1) ) Letter.objects.create(user=self.user, type="VAULT", status="SEALED", unlock_at=datetime.now(UTC)) Letter.objects.create( user=self.user, type="VAULT", status="SEALED", unlock_at=datetime.now(UTC) - timedelta(seconds=1) ) Letter.objects.create( user=self.user, type="VAULT", status="SEALED", unlock_at=datetime.now(UTC) - timedelta(hours=1), notified_at=datetime.now(UTC) - timedelta(minutes=59), ) Letter.objects.create( user=self.user, type="VAULT", status="SEALED", unlock_at=datetime.now(UTC) + timedelta(seconds=1) ) Letter.objects.create( user=self.user, type="KEPT", status="SEALED", ) unlocked_letters = get_vault_letters_to_notify() self.assertEqual(len(unlocked_letters), 2) def test_notify_unlocked_letter(self): """ Test that the task successfully notifies the user via email and updates the database field. """ from letters.tasks import notify_unlocked_letter letter_to_notify1 = Letter.objects.create( user=self.user, type="VAULT", status="SEALED", unlock_at=datetime.now(UTC), notified_at=None ) with patch("tasks.send_mail") as mock_send_mail: notify_unlocked_letter(letter_to_notify1) mock_send_mail.assert_called_with( subject=ANY, message=ANY, from_email=settings.FROM_EMAIL, recipient_list=[self.user.email], fail_silently=False, ) self.assertIsNotNone(letter_to_notify1.notified_at) letter_to_notify2 = Letter.objects.create( user=self.user, type="VAULT", status="SEALED", unlock_at=datetime.now(UTC), notified_at=None ) with patch("tasks.send_mail") as mock_send_mail: mock_send_mail.side_effect = Exception() notify_unlocked_letter(letter_to_notify2) self.assertIsNone(letter_to_notify2.notified_at)