From 2d912cef4284a207a5e526e14087e30dba425614 Mon Sep 17 00:00:00 2001
From: Ben Thorner <ben.thorner@digital.cabinet-office.gov.uk>
Date: Tue, 24 Aug 2021 15:11:00 +0100
Subject: [PATCH] Support area aggregation for custom polygons

---
 app/broadcast_areas/utils.py            | 16 ++++++++
 tests/app/broadcast_areas/test_utils.py | 52 +++++++++++++++++++++++++
 2 files changed, 68 insertions(+)

diff --git a/app/broadcast_areas/utils.py b/app/broadcast_areas/utils.py
index fc50d95c6..26da631cb 100644
--- a/app/broadcast_areas/utils.py
+++ b/app/broadcast_areas/utils.py
@@ -1,8 +1,24 @@
+from app.broadcast_areas.models import CustomBroadcastArea
+
+
 def aggregate_areas(areas):
+    areas = _convert_custom_areas_to_wards(areas)
     areas = _aggregate_wards_by_local_authority(areas)
     return areas
 
 
+def _convert_custom_areas_to_wards(areas):
+    results = set()
+
+    for area in areas:
+        if type(area) == CustomBroadcastArea:
+            results |= set(area.overlapping_electoral_wards)
+        else:
+            results |= {area}
+
+    return results
+
+
 def _aggregate_wards_by_local_authority(areas):
     return {
         area.parent if area.id.startswith('wd20-')
diff --git a/tests/app/broadcast_areas/test_utils.py b/tests/app/broadcast_areas/test_utils.py
index d9db48c0e..74e940fcc 100644
--- a/tests/app/broadcast_areas/test_utils.py
+++ b/tests/app/broadcast_areas/test_utils.py
@@ -69,3 +69,55 @@ def test_aggregate_areas(
     assert sorted(
         area.name for area in aggregate_areas(broadcast_message.areas)
     ) == expected_area_names
+
+
+@pytest.mark.parametrize(('simple_polygons', 'expected_area_names'), [
+    (
+        [SKYE], [
+            # Area covers a single unitary authority
+            'Highland',
+        ]
+    ),
+    (
+        [BRISTOL], [
+            # Area covers a single unitary authority
+            'Bristol, City of',
+        ]
+    ),
+    (
+        [CHELTENHAM], [
+            # Area covers three lower-tier authorities
+            # in Gloucestershire (upper tier authority)
+            'Cheltenham',
+            'Cotswold',
+            'Tewkesbury',
+        ]
+    ),
+    (
+        [BURFORD], [
+            # Area covers two lower-tier authorities
+            # in Gloucestershire (upper tier authority)
+            'Cotswold',
+            'West Oxfordshire',
+        ],
+    ),
+    (
+        [SANTA_A], [
+            # Does not overlap with the UK
+        ],
+    ),
+])
+def test_aggregate_areas_for_custom_polygons(
+    simple_polygons,
+    expected_area_names,
+):
+    broadcast_message = BroadcastMessage(
+        broadcast_message_json(
+            area_ids=['derived from polygons'],
+            simple_polygons=simple_polygons
+        )
+    )
+
+    assert sorted(
+        area.name for area in aggregate_areas(broadcast_message.areas)
+    ) == expected_area_names