Fix batch retreat with memberships

retirement/tests/tests_viewset_Retreat.py

@@ -1,5 +1,8 @@
 import json
-from datetime import datetime, timedelta
+from datetime import (
+    datetime,
+    timedelta,
+)
 
 import pytz
 import responses
@@ -13,10 +16,19 @@
 from rest_framework import status
 from rest_framework.test import APIClient
 
-from blitz_api.factories import AdminFactory, UserFactory
+from blitz_api.factories import (
+    AdminFactory,
+    UserFactory,
+)
+from blitz_api.models import AcademicLevel
 from blitz_api.testing_tools import CustomAPITestCase
+from store.models import Membership
 
-from ..models import Retreat, RetreatType, RetreatDate
+from retirement.models import (
+    Retreat,
+    RetreatType,
+    RetreatDate,
+)
 
 User = get_user_model()
 
@@ -186,6 +198,17 @@ def setUp(self):
         )
         self.retreat_hidden.activate()
 
+        self.academic_level = AcademicLevel.objects.create(
+            name="University"
+        )
+        self.membership = Membership.objects.create(
+            name="basic_membership",
+            details="1-Year student membership",
+            price=50,
+            available=True,
+            duration=timedelta(days=365),
+        )
+
     @override_settings(
         EXTERNAL_SCHEDULER={
             'URL': "http://example.com",
@@ -462,6 +485,12 @@ def test_create_batch_retreat(self):
                 'retreat:retreattype-detail',
                 args=[self.retreatType.id]
             ),
+            'exclusive_memberships': [
+                reverse(
+                    'membership-detail',
+                    args=[self.membership.id]
+                )
+            ],
             'bulk_start_time': '2021-03-01T08:00:00Z',
             'bulk_end_time': '2021-03-30T12:00:00Z',
             'weekdays': [0, 1, 2]
@@ -506,6 +535,15 @@ def test_create_batch_retreat(self):
 
         for item in content:
             self.assertEqual(item['name'][-2:], 'AM', item)
+            self.assertEqual(
+                item['exclusive_memberships'],
+                [
+                    'http://testserver' + reverse(
+                        'membership-detail',
+                        args=[self.membership.id]
+                    )
+                ],
+            )
             self.check_attributes(item, attributes)
 
     @override_settings(

retirement/views.py

@@ -214,6 +214,7 @@ def batch_create(self, request):
         validated_data.pop('bulk_start_time')
         validated_data.pop('bulk_end_time')
         validated_data.pop('weekdays')
+        memberships = validated_data.pop('exclusive_memberships')
 
         for start, end in zip(retreat_start_dates, retreat_end_dates):
             if tz.localize(start).hour < 12:
@@ -224,6 +225,7 @@ def batch_create(self, request):
                 "Bloc %d %b"
             ) + ' ' + suffix
             new_retreat = Retreat.objects.create(**validated_data)
+            new_retreat.exclusive_memberships.set(memberships)
             RetreatDate.objects.create(
                 retreat=new_retreat,
                 start_time=tz.localize(start),