From 1693473fd22a14238262a9655fc129276aa6fe68 Mon Sep 17 00:00:00 2001
From: Adam Kelly <adam@cthulahoops.org>
Date: Fri, 13 Mar 2020 13:38:09 +0000
Subject: [PATCH] Use Enum and constraint for valid notification types.

---
 .../migrations/0016_auto_20200313_1337.py     | 22 ++++++++++++++++
 fedireads/models/status.py                    | 25 ++++++++-----------
 2 files changed, 33 insertions(+), 14 deletions(-)
 create mode 100644 fedireads/migrations/0016_auto_20200313_1337.py

diff --git a/fedireads/migrations/0016_auto_20200313_1337.py b/fedireads/migrations/0016_auto_20200313_1337.py
new file mode 100644
index 00000000..4f48de45
--- /dev/null
+++ b/fedireads/migrations/0016_auto_20200313_1337.py
@@ -0,0 +1,22 @@
+# Generated by Django 3.0.3 on 2020-03-13 13:37
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('fedireads', '0015_auto_20200311_1212'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='notification',
+            name='notification_type',
+            field=models.CharField(choices=[('FAVORITE', 'Favorite'), ('REPLY', 'Reply'), ('TAG', 'Tag'), ('FOLLOW', 'Follow'), ('FOLLOW_REQUEST', 'Follow Request')], max_length=255),
+        ),
+        migrations.AddConstraint(
+            model_name='notification',
+            constraint=models.CheckConstraint(check=models.Q(notification_type__in=['FAVORITE', 'REPLY', 'TAG', 'FOLLOW', 'FOLLOW_REQUEST']), name='notification_type_valid'),
+        ),
+    ]
diff --git a/fedireads/models/status.py b/fedireads/models/status.py
index b2e83895..228f48cc 100644
--- a/fedireads/models/status.py
+++ b/fedireads/models/status.py
@@ -88,6 +88,9 @@ class Tag(FedireadsModel):
         unique_together = ('user', 'book', 'name')
 
 
+NotificationType = models.TextChoices(
+    'NotificationType', 'FAVORITE REPLY TAG FOLLOW FOLLOW_REQUEST')
+
 class Notification(FedireadsModel):
     ''' you've been tagged, liked, followed, etc '''
     user = models.ForeignKey('User', on_delete=models.PROTECT)
@@ -99,18 +102,12 @@ class Notification(FedireadsModel):
     related_status = models.ForeignKey(
         'Status', on_delete=models.PROTECT, null=True)
     read = models.BooleanField(default=False)
-    notification_type = models.CharField(max_length=255)
-
-    def save(self, *args, **kwargs):
-        # TODO: there's probably a real way to do enums
-        types = [
-            'FAVORITE',
-            'REPLY',
-            'TAG',
-            'FOLLOW',
-            'FOLLOW_REQUEST'
+    notification_type = models.CharField(
+        max_length=255, choices=NotificationType.choices)
+    class Meta:
+        constraints = [
+            models.CheckConstraint(
+                check=models.Q(notification_type__in=NotificationType.values),
+                name="notification_type_valid",
+            )
         ]
-        if not self.notification_type in types:
-            raise ValueError('Invalid notitication type')
-        super().save(*args, **kwargs)
-