ensure you can get broadcast_areas from a list as well as from *args

This commit is contained in:
Leo Hemsted
2020-07-09 19:10:51 +01:00
committed by Toby Lorne
parent 078f1dd8d3
commit b3ed21252d
2 changed files with 18 additions and 0 deletions

View File

@@ -147,6 +147,10 @@ class BroadcastAreaLibraries(SerialisedModelCollection, GetItemByIdMixin):
yield area
def get_areas(self, *area_ids):
# allow people to call `get_areas('a', 'b') or get_areas(['a', 'b'])`
if len(area_ids) == 1 and isinstance(area_ids[0], list):
area_ids = area_ids[0]
return list(itertools.chain(*(
[area for area in self.all_areas if area.id == area_id]
for area_id in area_ids

View File

@@ -100,6 +100,16 @@ def test_get_names_of_areas():
]
def test_get_areas_accepts_lists():
areas_from_list = broadcast_area_libraries.get_areas(
['wales', 'vale-of-glamorgan', 'england', 'essex']
)
areas_from_args = broadcast_area_libraries.get_areas(
'wales', 'vale-of-glamorgan', 'england', 'essex',
)
assert areas_from_args == areas_from_list
def test_has_polygons():
assert len(
@@ -114,6 +124,10 @@ def test_has_polygons():
broadcast_area_libraries.get_polygons_for_areas_long_lat('england', 'scotland')
) == 35 + 195 == 230
assert len(
broadcast_area_libraries.get_polygons_for_areas_long_lat(['england', 'scotland'])
) == 35 + 195 == 230
assert broadcast_area_libraries.get_polygons_for_areas_lat_long('england')[0][0] == [
55.811085, -2.034358 # https://goo.gl/maps/wsf2LUWzYinwydMk8
]