From bae94ef8294b86b6da3d2fa25180fe5a9dda88ed Mon Sep 17 00:00:00 2001
From: Erik Harding <erik@praekelt.org>
Date: Mon, 28 Aug 2023 15:00:42 +0200
Subject: [PATCH] Handle survey group not being set

---
 yal/surveys/location.py            |  4 ++++
 yal/tests/surveys/test_location.py | 17 +++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/yal/surveys/location.py b/yal/surveys/location.py
index a06d889f..384e959b 100644
--- a/yal/surveys/location.py
+++ b/yal/surveys/location.py
@@ -23,6 +23,10 @@ async def state_location_introduction(self):
             "ejaf_location_survey_status", "not_invited"
         )
         survey_group = self.user.metadata.get("ejaf_location_survey_group")
+
+        if not survey_group:
+            return await self.go_to_state("state_location_not_invited")
+
         self.save_answer("ejaf_location_survey_group", survey_group)
 
         error, group_count = await rapidpro.get_group_membership_count(
diff --git a/yal/tests/surveys/test_location.py b/yal/tests/surveys/test_location.py
index e776f57c..7210ad1e 100644
--- a/yal/tests/surveys/test_location.py
+++ b/yal/tests/surveys/test_location.py
@@ -21,10 +21,13 @@ def get_rapidpro_contact(urn):
         status = "completed"
     if "27820001003" in urn:
         status = None
+        survey_group = None
     if "27820001004" in urn:
         survey_group = "2"
     if "27820001005" in urn:
         status = "invalid_province"
+    if "27820001006" in urn:
+        status = None
     return {
         "fields": {
             "ejaf_location_survey_status": status,
@@ -137,6 +140,20 @@ async def test_state_location_introduction_not_invited(tester: AppTester):
     )
 
 
+@pytest.mark.asyncio
+async def test_state_location_introduction_not_invited_with_group(tester: AppTester):
+    tester.setup_user_address("27820001006")
+    tester.setup_state("state_location_introduction")
+    await tester.user_input(session=Message.SESSION_EVENT.NEW)
+
+    tester.assert_state("state_location_not_invited")
+
+    tester.assert_message(
+        "Unfortunately it looks like we already have enough people answering this "
+        "survey, but thank you for your interest."
+    )
+
+
 @pytest.mark.asyncio
 async def test_state_location_introduction_group_max(tester: AppTester):
     tester.setup_user_address("27820001004")