Add text only error to FreeText state

vaccine/states.py

@@ -254,6 +254,23 @@ async def display(self, message: Message):
         return self.app.send_message(text, helper_metadata=helper_metadata)
 
 
+class WhatsAppTextOnlyFreetext(FreeText):
+    def __init__(
+        self,
+        *args,
+        text_only_error: Optional[str] = None,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.text_only_error = text_only_error
+
+    async def process_message(self, message: Message):
+        if self.text_only_error is not None and message.transport_metadata:
+            if message.transport_metadata.get("message", {}).get("type") != "text":
+                return self.app.send_message(self.text_only_error)
+        return await super().process_message(message)
+
+
 class BaseWhatsAppChoiceState(ChoiceState):
     """
     Base class for all whatsapp choice states.

yal/surveys/location.py

@@ -4,9 +4,9 @@
 from vaccine.states import (
     Choice,
     EndState,
-    FreeText,
     WhatsAppButtonState,
     WhatsAppListState,
+    WhatsAppTextOnlyFreetext,
 )
 from vaccine.validators import nonempty_validator
 from yal import config, contentrepo, rapidpro
@@ -258,11 +258,12 @@ async def state_location_name_city(self):
                 ]
             )
         )
-        return FreeText(
+        return WhatsAppTextOnlyFreetext(
             self,
             question=question,
             next="state_location_area_type",
             check=nonempty_validator(question),
+            text_only_error=question,
         )
 
     async def state_location_area_type(self):

yal/tests/surveys/test_location.py

@@ -308,6 +308,32 @@ async def test_state_location_name_city(tester: AppTester):
     )
 
 
+@pytest.mark.asyncio
+async def test_state_location_name_city_invalid(tester: AppTester):
+    tester.setup_state("state_location_name_city")
+    await tester.user_input(
+        "",
+        transport_metadata={
+            "message": {
+                "type": "image",
+                "image": {"id": "img1", "mime_type": "image/jpeg"},
+            }
+        },
+    )
+
+    tester.assert_state("state_location_name_city")
+    tester.assert_message(
+        "\n".join(
+            [
+                "*What is the name of the city or town you live in or live closest "
+                "to?*",
+                "",
+                "Please *TYPE* in the name of the city or town.",
+            ]
+        )
+    )
+
+
 @pytest.mark.asyncio
 async def test_state_location_area_type(tester: AppTester):
     tester.setup_state("state_location_name_city")