Move test_schedule_bin into the schedules management command base test class

openedx/core/djangoapps/schedules/management/commands/tests/test_send_recurring_nudge.py

@@ -54,36 +54,6 @@ class TestSendRecurringNudge(ScheduleBaseEmailTestBase):
     tested_command = nudge.Command
     expected_offsets = (-3, -10)
 
-
-    @ddt.data(1, 10, 100)
-    @patch.object(tasks, 'ace')
-    @patch.object(tested_task, 'async_send_task')
-    def test_schedule_bin(self, schedule_count, mock_schedule_send, mock_ace):
-        schedules = [
-            ScheduleFactory.create(
-                start=datetime.datetime(2017, 8, 3, 18, 44, 30, tzinfo=pytz.UTC),
-                enrollment__course__id=CourseLocator('edX', 'toy', 'Bin')
-            ) for i in range(schedule_count)
-        ]
-
-        bins_in_use = frozenset((s.enrollment.user.id % resolvers.RECURRING_NUDGE_NUM_BINS) for s in schedules)
-
-        test_datetime = datetime.datetime(2017, 8, 3, 18, tzinfo=pytz.UTC)
-        test_datetime_str = serialize(test_datetime)
-        for b in range(resolvers.RECURRING_NUDGE_NUM_BINS):
-            expected_queries = NUM_QUERIES_NO_MATCHING_SCHEDULES + NUM_QUERIES_NO_ORG_LIST
-            if b in bins_in_use:
-                # to fetch course modes for valid schedules
-                expected_queries += NUM_COURSE_MODES_QUERIES
-
-            with self.assertNumQueries(expected_queries, table_blacklist=WAFFLE_TABLES):
-
-                self.tested_task.apply(kwargs=dict(
-                    site_id=self.site_config.site.id, target_day_str=test_datetime_str, day_offset=-3, bin_num=b,
-                ))
-        self.assertEqual(mock_schedule_send.apply_async.call_count, schedule_count)
-        self.assertFalse(mock_ace.send.called)
-
     @patch.object(tested_task, 'async_send_task')
     def test_no_course_overview(self, mock_schedule_send):
         schedule = ScheduleFactory.create(

openedx/core/djangoapps/schedules/management/commands/tests/test_send_upgrade_reminder.py

@@ -8,7 +8,6 @@ import ddt
 import pytz
 from django.conf import settings
 from edx_ace import Message
-from freezegun import freeze_time
 from edx_ace.channel import ChannelType
 from edx_ace.test_utils import StubPolicy, patch_channels, patch_policies
 from edx_ace.utils.date import serialize
@@ -18,7 +17,6 @@ from opaque_keys.edx.locator import CourseLocator
 
 from course_modes.models import CourseMode
 from course_modes.tests.factories import CourseModeFactory
-from openedx.core.djangoapps.content.course_overviews.models import CourseOverview
 from openedx.core.djangoapps.schedules import resolvers, tasks
 from openedx.core.djangoapps.schedules.management.commands import send_upgrade_reminder as reminder
 from openedx.core.djangoapps.schedules.management.commands.tests.tools import ScheduleBaseEmailTestBase
@@ -27,8 +25,6 @@ from openedx.core.djangoapps.site_configuration.tests.factories import SiteConfi
 from openedx.core.djangoapps.waffle_utils.testutils import WAFFLE_TABLES
 from openedx.core.djangolib.testing.utils import skip_unless_lms
 from student.tests.factories import UserFactory
-from xmodule.modulestore.tests.django_utils import SharedModuleStoreTestCase
-from xmodule.modulestore.tests.factories import CourseFactory
 
 
 SITE_QUERY = 1
@@ -58,26 +54,14 @@ LOG = logging.getLogger(__name__)
 @skip_unless_lms
 @skipUnless('openedx.core.djangoapps.schedules.apps.SchedulesConfig' in settings.INSTALLED_APPS,
             "Can't test schedules if the app isn't installed")
-@freeze_time('2017-08-01 00:00:00', tz_offset=0, tick=True)
-class TestUpgradeReminder(ScheduleBaseEmailTestBase, SharedModuleStoreTestCase):
+class TestUpgradeReminder(ScheduleBaseEmailTestBase):
     __test__ = True
 
     tested_task = tasks.ScheduleUpgradeReminder
     tested_command = reminder.Command
     expected_offsets = (2,)
 
-    @classmethod
-    def setUpClass(cls):
-        super(TestUpgradeReminder, cls).setUpClass()
-
-        cls.course = CourseFactory.create(
-            org='edX',
-            number='test',
-            display_name='Test Course',
-            self_paced=True,
-            start=datetime.datetime.now(pytz.UTC) - datetime.timedelta(days=30),
-        )
-        cls.course_overview = CourseOverview.get_from_id(cls.course.id)
+    has_course_queries = True
 
     def setUp(self):
         super(TestUpgradeReminder, self).setUp()
@@ -88,50 +72,6 @@ class TestUpgradeReminder(ScheduleBaseEmailTestBase, SharedModuleStoreTestCase):
             expiration_datetime=datetime.datetime.now(pytz.UTC) + datetime.timedelta(days=30),
         )
 
-    @ddt.data(1, 10, 100)
-    @patch.object(tasks, 'ace')
-    @patch.object(tested_task, 'async_send_task')
-    def test_schedule_bin(self, schedule_count, mock_schedule_send, mock_ace):
-        upgrade_deadline = datetime.datetime.now(pytz.UTC) + datetime.timedelta(days=2)
-        schedules = [
-            ScheduleFactory.create(
-                upgrade_deadline=upgrade_deadline,
-                enrollment__course=self.course_overview,
-            ) for i in range(schedule_count)
-        ]
-
-        bins_in_use = frozenset((self._calculate_bin_for_user(s.enrollment.user)) for s in schedules)
-        is_first_match = True
-        course_switch_queries = len(set(s.enrollment.course.id for s in schedules))
-        org_switch_queries = len(set(s.enrollment.course.id.org for s in schedules))
-        test_datetime = upgrade_deadline
-        test_datetime_str = serialize(test_datetime)
-
-        for b in range(resolvers.UPGRADE_REMINDER_NUM_BINS):
-            LOG.debug('Running bin %d', b)
-            expected_queries = NUM_QUERIES_NO_MATCHING_SCHEDULES
-            if b in bins_in_use:
-                if is_first_match:
-                    expected_queries = (
-                        # Since this is the first match, we need to cache all of the config models, so we run a query
-                        # for each of those...
-                        NUM_QUERIES_FIRST_MATCH
-                        + course_switch_queries + org_switch_queries
-                    )
-                    is_first_match = False
-                else:
-                    expected_queries = NUM_QUERIES_WITH_MATCHES
-
-            expected_queries += NUM_QUERIES_NO_ORG_LIST
-
-            with self.assertNumQueries(expected_queries, table_blacklist=WAFFLE_TABLES):
-                self.tested_task.apply(kwargs=dict(
-                    site_id=self.site_config.site.id, target_day_str=test_datetime_str, day_offset=2, bin_num=b,
-                ))
-
-        self.assertEqual(mock_schedule_send.apply_async.call_count, schedule_count)
-        self.assertFalse(mock_ace.send.called)
-
     @patch.object(tested_task, 'async_send_task')
     def test_no_course_overview(self, mock_schedule_send):
 

openedx/core/djangoapps/schedules/management/commands/tests/tools.py

 import datetime
+import ddt
+import logging
 
-from edx_ace.utils.date import serialize
+from freezegun import freeze_time
 from mock import patch
 import pytz
 
 from courseware.models import DynamicUpgradeDeadlineConfiguration
-from openedx.core.djangolib.testing.utils import CacheIsolationTestCase, FilteredQueryCountMixin
+from edx_ace.utils.date import serialize
+from opaque_keys.edx.locator import CourseLocator
+from openedx.core.djangoapps.content.course_overviews.models import CourseOverview
 from openedx.core.djangoapps.site_configuration.tests.factories import SiteConfigurationFactory, SiteFactory
-from openedx.core.djangoapps.schedules import tasks
-from openedx.core.djangoapps.schedules.tests.factories import ScheduleConfigFactory
+from openedx.core.djangoapps.schedules import resolvers, tasks
+from openedx.core.djangoapps.schedules.tests.factories import ScheduleConfigFactory, ScheduleFactory
+from openedx.core.djangoapps.waffle_utils.testutils import WAFFLE_TABLES
+from openedx.core.djangolib.testing.utils import CacheIsolationTestCase, FilteredQueryCountMixin
+from xmodule.modulestore.tests.django_utils import SharedModuleStoreTestCase
+from xmodule.modulestore.tests.factories import CourseFactory
+
+
+SITE_QUERY = 1
+SCHEDULES_QUERY = 1
+COURSE_MODES_QUERY = 1
+GLOBAL_DEADLINE_SWITCH_QUERY = 1
+COMMERCE_CONFIG_QUERY = 1
+NUM_QUERIES_NO_ORG_LIST = 1
 
+NUM_QUERIES_NO_MATCHING_SCHEDULES = SITE_QUERY + SCHEDULES_QUERY
 
+NUM_QUERIES_WITH_MATCHES = (
+    NUM_QUERIES_NO_MATCHING_SCHEDULES +
+    COURSE_MODES_QUERY
+)
 
-class ScheduleBaseEmailTestBase(FilteredQueryCountMixin, CacheIsolationTestCase):
+NUM_QUERIES_FIRST_MATCH = (
+    NUM_QUERIES_WITH_MATCHES
+    + GLOBAL_DEADLINE_SWITCH_QUERY
+    + COMMERCE_CONFIG_QUERY
+)
+
+LOG = logging.getLogger(__name__)
+
+
+@ddt.ddt
+@freeze_time('2017-08-01 00:00:00', tz_offset=0, tick=True)
+class ScheduleBaseEmailTestBase(SharedModuleStoreTestCase):
 
     __test__ = False
 
     ENABLED_CACHES = ['default']
 
+    has_course_queries = False
+
+    @classmethod
+    def setUpClass(cls):
+        super(ScheduleBaseEmailTestBase, cls).setUpClass()
+
+        cls.course = CourseFactory.create(
+            org='edX',
+            number='test',
+            display_name='Test Course',
+            self_paced=True,
+            start=datetime.datetime.now(pytz.UTC) - datetime.timedelta(days=30),
+        )
+        cls.course_overview = CourseOverview.get_from_id(cls.course.id)
+
     def setUp(self):
         super(ScheduleBaseEmailTestBase, self).setUp()
 
@@ -60,3 +107,52 @@ class ScheduleBaseEmailTestBase(FilteredQueryCountMixin, CacheIsolationTestCase)
                 retry=False,
             )
             self.assertFalse(mock_ace.send.called)
+
+    @ddt.data(1, 10, 100)
+    @patch.object(tasks, 'ace')
+    @patch.object(resolvers, 'set_custom_metric')
+    def test_schedule_bin(self, schedule_count, mock_metric, mock_ace):
+        with patch.object(self.tested_task, 'async_send_task') as mock_schedule_send:
+            current_day, offset, target_day = self._get_dates()
+            schedules = [
+                ScheduleFactory.create(
+                    start=target_day,
+                    upgrade_deadline=target_day,
+                    enrollment__course__self_paced=True,
+                ) for _ in range(schedule_count)
+            ]
+
+            bins_in_use = frozenset((self._calculate_bin_for_user(s.enrollment.user)) for s in schedules)
+            is_first_match = True
+            course_queries = len(set(s.enrollment.course.id for s in schedules)) if self.has_course_queries else 0
+            target_day_str = serialize(target_day)
+
+            for b in range(self.tested_task.num_bins):
+                LOG.debug('Running bin %d', b)
+                expected_queries = NUM_QUERIES_NO_MATCHING_SCHEDULES
+                if b in bins_in_use:
+                    if is_first_match:
+                        expected_queries = (
+                            # Since this is the first match, we need to cache all of the config models, so we run a
+                            # query for each of those...
+                            NUM_QUERIES_FIRST_MATCH + course_queries
+                        )
+                        is_first_match = False
+                    else:
+                        expected_queries = NUM_QUERIES_WITH_MATCHES
+
+                expected_queries += NUM_QUERIES_NO_ORG_LIST
+
+                with self.assertNumQueries(expected_queries, table_blacklist=WAFFLE_TABLES):
+                    self.tested_task.apply(kwargs=dict(
+                        site_id=self.site_config.site.id, target_day_str=target_day_str, day_offset=offset, bin_num=b,
+                    ))
+
+                num_schedules = mock_metric.call_args[0][1]
+                if b in bins_in_use:
+                    self.assertGreater(num_schedules, 0)
+                else:
+                    self.assertEqual(num_schedules, 0)
+
+            self.assertEqual(mock_schedule_send.apply_async.call_count, schedule_count)
+            self.assertFalse(mock_ace.send.called)