From 2e04667faded1124cbe16074a857a38b36c29405 Mon Sep 17 00:00:00 2001
From: Jumana B <jzbahrai@users.noreply.github.com>
Date: Tue, 9 Jul 2024 14:55:00 -0400
Subject: [PATCH] updated the template category model (#2211)

---
 app/models.py                                 |  7 +++
 tests/app/dao/test_template_categories_dao.py | 51 ++++++++++++++-----
 2 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/app/models.py b/app/models.py
index 2d0e665c6a..81e9d33f26 100644
--- a/app/models.py
+++ b/app/models.py
@@ -65,6 +65,11 @@
 COMPLAINT_CALLBACK_TYPE = "complaint"
 SERVICE_CALLBACK_TYPES = [DELIVERY_STATUS_CALLBACK_TYPE, COMPLAINT_CALLBACK_TYPE]
 
+SHORT_CODE = "short_code"
+LONG_CODE = "long_code"
+
+sms_sending_vehicles = db.Enum(*[SHORT_CODE, LONG_CODE], name="sms_sending_vehicles")
+
 
 def filter_null_value_fields(obj):
     return dict(filter(lambda x: x[1] is not None, obj.items()))
@@ -1046,6 +1051,7 @@ class TemplateCategory(BaseModel):
     hidden = db.Column(db.Boolean, nullable=False, default=False)
     created_at = db.Column(db.DateTime, nullable=False, default=datetime.datetime.utcnow)
     updated_at = db.Column(db.DateTime, onupdate=datetime.datetime.utcnow)
+    sms_sending_vehicle = db.Column(sms_sending_vehicles, nullable=False, default="long_code")
 
     def serialize(self):
         return {
@@ -1059,6 +1065,7 @@ def serialize(self):
             "hidden": self.hidden,
             "created_at": self.created_at,
             "updated_at": self.updated_at,
+            "sms_sending_vehicle": self.sms_sending_vehicle,
         }
 
     @classmethod
diff --git a/tests/app/dao/test_template_categories_dao.py b/tests/app/dao/test_template_categories_dao.py
index 73832a2a43..b9d242a516 100644
--- a/tests/app/dao/test_template_categories_dao.py
+++ b/tests/app/dao/test_template_categories_dao.py
@@ -14,22 +14,45 @@
 from tests.app.conftest import create_sample_template
 
 
-def test_create_template_category(notify_db_session):
-    data = {
-        "name_en": "english",
-        "name_fr": "french",
-        "description_en": "english description",
-        "description_fr": "french description",
-        "sms_process_type": NORMAL,
-        "email_process_type": NORMAL,
-        "hidden": False,
-    }
+class TestCreateTemplateCategory:
+    def test_create_template_category(self, notify_db_session):
+        data = {
+            "name_en": "english",
+            "name_fr": "french",
+            "description_en": "english description",
+            "description_fr": "french description",
+            "sms_process_type": NORMAL,
+            "email_process_type": NORMAL,
+            "hidden": False,
+            "sms_sending_vehicle": "short_code",
+        }
+
+        template_category = TemplateCategory(**data)
+        dao_create_template_category(template_category)
 
-    template_category = TemplateCategory(**data)
-    dao_create_template_category(template_category)
+        temp_cat = dao_get_all_template_categories()
+        assert TemplateCategory.query.count() == 1
+        assert len(temp_cat) == 1
+        assert temp_cat[0].sms_sending_vehicle == "short_code"
+
+    def test_create_template_category_with_no_sms_sending_vehicle(self, notify_db_session):
+        data = {
+            "name_en": "english",
+            "name_fr": "french",
+            "description_en": "english description",
+            "description_fr": "french description",
+            "sms_process_type": NORMAL,
+            "email_process_type": NORMAL,
+            "hidden": False,
+        }
+
+        template_category = TemplateCategory(**data)
+        dao_create_template_category(template_category)
 
-    assert TemplateCategory.query.count() == 1
-    assert len(dao_get_all_template_categories()) == 1
+        temp_cat = dao_get_all_template_categories()
+        assert TemplateCategory.query.count() == 1
+        assert len(temp_cat) == 1
+        assert temp_cat[0].sms_sending_vehicle == "long_code"  # default value
 
 
 @pytest.mark.parametrize(