Add OrgWaffleFlag to waffle_utils

(cherry picked from commit deb2d3a542d3fe050a8e29eccb11a0826a908c44)

openedx/core/djangoapps/waffle_utils/__init__.py

@@ -360,3 +360,61 @@ class CourseWaffleFlag(WaffleFlag):
             check_before_waffle_callback=self._get_course_override_callback(course_key),
             flag_undefined_default=self.flag_undefined_default
         )
+
+
+class OrgWaffleFlag(WaffleFlag):
+    """
+    Represents a single waffle flag that can be forced on/off for an organization.
+
+    Uses a cached waffle namespace.
+    """
+
+    def _get_org_override_callback(self, org_key):
+        """
+        Returns a function to use as the check_before_waffle_callback.
+
+        Arguments:
+            org_key (String): The org component of the CourseKey to check for override
+            before checking waffle.
+        """
+        def org_override_callback(namespaced_flag_name):
+            """
+            Returns True/False if the flag was forced on or off for the provided
+            org.  Returns None if the flag was not overridden.
+
+            Note: Has side effect of caching the override value.
+
+            Arguments:
+                namespaced_flag_name (String): A namespaced version of the flag
+                    to check.
+            """
+            # Import is placed here to avoid model import at project startup.
+            from .models import WaffleFlagOrgOverrideModel
+            cache_key = u'{}.{}'.format(namespaced_flag_name, unicode(org_key))
+            force_override = self.waffle_namespace._cached_flags.get(cache_key)
+
+            if force_override is None:
+                force_override = WaffleFlagOrgOverrideModel.override_value(namespaced_flag_name, org_key)
+                self.waffle_namespace._cached_flags[cache_key] = force_override
+
+            if force_override == WaffleFlagOrgOverrideModel.ALL_CHOICES.on:
+                return True
+            if force_override == WaffleFlagOrgOverrideModel.ALL_CHOICES.off:
+                return False
+            return None
+
+        return org_override_callback
+
+    def is_enabled(self, org_key=None):
+        """
+        Returns whether or not the flag is enabled.
+
+        Arguments:
+            org_key (String): The org component of the CourseKey to check for override before
+            checking waffle.
+        """
+        return self.waffle_namespace.is_flag_active(
+            self.flag_name,
+            check_before_waffle_callback=self._get_org_override_callback(org_key),
+            flag_undefined_default=self.flag_undefined_default
+        )

openedx/core/djangoapps/waffle_utils/admin.py

@@ -3,11 +3,11 @@ Django admin page for waffle utils models
 """
 from django.contrib import admin
 
-from config_models.admin import ConfigurationModelAdmin, KeyedConfigurationModelAdmin
+from config_models.admin import KeyedConfigurationModelAdmin
 
 
-from .forms import WaffleFlagCourseOverrideAdminForm
-from .models import WaffleFlagCourseOverrideModel
+from .forms import WaffleFlagCourseOverrideAdminForm, WaffleFlagOrgOverrideAdminForm
+from .models import WaffleFlagCourseOverrideModel, WaffleFlagOrgOverrideModel
 
 
 class WaffleFlagCourseOverrideAdmin(KeyedConfigurationModelAdmin):
@@ -26,4 +26,23 @@ class WaffleFlagCourseOverrideAdmin(KeyedConfigurationModelAdmin):
         }),
     )
 
+
+class WaffleFlagOrgOverrideAdmin(KeyedConfigurationModelAdmin):
+    """
+    Admin for course override of waffle flags.
+
+    Includes search by org_id and waffle_flag.
+
+    """
+    form = WaffleFlagOrgOverrideAdminForm
+    search_fields = ['waffle_flag', 'org_id']
+    fieldsets = (
+        (None, {
+            'fields': ('waffle_flag', 'org_id', 'override_choice', 'enabled'),
+            'description': 'Enter a valid org id and an existing waffle flag. The waffle flag name is not validated.'
+        }),
+    )
+
+
 admin.site.register(WaffleFlagCourseOverrideModel, WaffleFlagCourseOverrideAdmin)
+admin.site.register(WaffleFlagOrgOverrideModel, WaffleFlagOrgOverrideAdmin)

openedx/core/djangoapps/waffle_utils/forms.py

@@ -5,7 +5,7 @@ from django import forms
 
 from openedx.core.lib.courses import clean_course_id
 
-from .models import WaffleFlagCourseOverrideModel
+from .models import WaffleFlagCourseOverrideModel, WaffleFlagOrgOverrideModel
 
 
 class WaffleFlagCourseOverrideAdminForm(forms.ModelForm):
@@ -33,3 +33,24 @@ class WaffleFlagCourseOverrideAdminForm(forms.ModelForm):
             raise forms.ValidationError(msg)
 
         return cleaned_flag.strip()
+
+
+class WaffleFlagOrgOverrideAdminForm(forms.ModelForm):
+    """
+    Input form for org override of waffle flags, allowing us to verify data.
+    """
+    class Meta(object):
+        model = WaffleFlagOrgOverrideModel
+        fields = '__all__'
+
+    def clean_waffle_flag(self):
+        """
+        Validate the waffle flag is an existing flag.
+        """
+        cleaned_flag = self.cleaned_data['waffle_flag']
+
+        if not cleaned_flag:
+            msg = u'Waffle flag must be supplied.'
+            raise forms.ValidationError(msg)
+
+        return cleaned_flag.strip()

openedx/core/djangoapps/waffle_utils/migrations/0002_waffleflagorgoverridemodel.py

+# -*- coding: utf-8 -*-
+from __future__ import unicode_literals
+
+from django.db import migrations, models
+import django.db.models.deletion
+from django.conf import settings
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ('waffle_utils', '0001_initial'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='WaffleFlagOrgOverrideModel',
+            fields=[
+                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
+                ('change_date', models.DateTimeField(auto_now_add=True, verbose_name='Change date')),
+                ('enabled', models.BooleanField(default=False, verbose_name='Enabled')),
+                ('waffle_flag', models.CharField(max_length=255, db_index=True)),
+                ('org_id', models.CharField(max_length=255, db_index=True)),
+                ('override_choice', models.CharField(default=b'on', max_length=3, choices=[(b'on', 'Force On'), (b'off', 'Force Off')])),
+                ('changed_by', models.ForeignKey(on_delete=django.db.models.deletion.PROTECT, editable=False, to=settings.AUTH_USER_MODEL, null=True, verbose_name='Changed by')),
+            ],
+            options={
+                'verbose_name': 'Waffle flag org override',
+                'verbose_name_plural': 'Waffle flag org overrides',
+            },
+        ),
+    ]

openedx/core/djangoapps/waffle_utils/models.py

@@ -59,3 +59,54 @@ class WaffleFlagCourseOverrideModel(ConfigurationModel):
         enabled_label = "Enabled" if self.enabled else "Not Enabled"
         # pylint: disable=no-member
         return u"Course '{}': Persistent Grades {}".format(self.course_id.to_deprecated_string(), enabled_label)
+
+
+class WaffleFlagOrgOverrideModel(ConfigurationModel):
+    """
+    Used to force a waffle flag on or off for an organization.
+    """
+    OVERRIDE_CHOICES = Choices(('on', _('Force On')), ('off', _('Force Off')))
+    ALL_CHOICES = OVERRIDE_CHOICES + Choices('unset')
+
+    KEY_FIELDS = ('waffle_flag', 'org_id')
+
+    # The org that these features are attached to.
+    waffle_flag = CharField(max_length=255, db_index=True)
+    org_id = CharField(max_length=255, db_index=True)
+    override_choice = CharField(choices=OVERRIDE_CHOICES, default=OVERRIDE_CHOICES.on, max_length=3)
+
+    @classmethod
+    @request_cached
+    def override_value(cls, waffle_flag, org_id):
+        """
+        Returns whether the waffle flag was overridden (on or off) for the
+        org, or is unset.
+
+        Arguments:
+            waffle_flag (String): The name of the flag.
+            org_id (String): The org id for which the flag may have
+                been overridden.
+
+        If the current config is not set or disabled for this waffle flag and
+            org id, returns ALL_CHOICES.unset.
+        Otherwise, returns ALL_CHOICES.on or ALL_CHOICES.off as configured for
+            the override_choice.
+
+        """
+        if not org_id or not waffle_flag:
+            return cls.ALL_CHOICES.unset
+
+        effective = cls.objects.filter(waffle_flag=waffle_flag, org_id=org_id).order_by('-change_date').first()
+        if effective and effective.enabled:
+            return effective.override_choice
+        return cls.ALL_CHOICES.unset
+
+    class Meta(object):
+        app_label = "waffle_utils"
+        verbose_name = 'Waffle flag org override'
+        verbose_name_plural = 'Waffle flag org overrides'
+
+    def __unicode__(self):
+        enabled_label = "Enabled" if self.enabled else "Not Enabled"
+        # pylint: disable=no-member
+        return u"Org '{}': Persistent Grades {}".format(self.org_id, enabled_label)

openedx/core/djangoapps/waffle_utils/tests/test_init.py

@@ -8,8 +8,8 @@ from opaque_keys.edx.keys import CourseKey
 from request_cache.middleware import RequestCache
 from waffle.testutils import override_flag
 
-from .. import CourseWaffleFlag, WaffleFlagNamespace
-from ..models import WaffleFlagCourseOverrideModel
+from .. import CourseWaffleFlag, OrgWaffleFlag, WaffleFlagNamespace
+from ..models import WaffleFlagCourseOverrideModel, WaffleFlagOrgOverrideModel
 
 
 @ddt.ddt
@@ -91,3 +91,84 @@ class TestCourseWaffleFlag(TestCase):
                 self.NAMESPACED_FLAG_NAME,
                 self.TEST_COURSE_KEY
             )
+
+
+@ddt.ddt
+class TestOrgWaffleFlag(TestCase):
+    """
+    Tests the OrgWaffleFlag.
+    """
+
+    NAMESPACE_NAME = 'test_namespace'
+    FLAG_NAME = 'test_flag'
+    NAMESPACED_FLAG_NAME = NAMESPACE_NAME + '.' + FLAG_NAME
+
+    TEST_ORG_KEY = 'edX'
+    TEST_ORG_2_KEY = 'edX2'
+    TEST_NAMESPACE = WaffleFlagNamespace(NAMESPACE_NAME)
+    TEST_ORG_FLAG = OrgWaffleFlag(TEST_NAMESPACE, FLAG_NAME)
+
+    @ddt.data(
+        {'org_override': WaffleFlagOrgOverrideModel.ALL_CHOICES.on, 'waffle_enabled': False, 'result': True},
+        {'org_override': WaffleFlagOrgOverrideModel.ALL_CHOICES.off, 'waffle_enabled': True, 'result': False},
+        {'org_override': WaffleFlagOrgOverrideModel.ALL_CHOICES.unset, 'waffle_enabled': True, 'result': True},
+        {'org_override': WaffleFlagOrgOverrideModel.ALL_CHOICES.unset, 'waffle_enabled': False, 'result': False},
+    )
+    def test_org_waffle_flag(self, data):
+        """
+        Tests various combinations of a flag being set in waffle and overridden
+        for an organization.
+        """
+        RequestCache.clear_request_cache()
+
+        with patch.object(WaffleFlagOrgOverrideModel, 'override_value', return_value=data['org_override']):
+            with override_flag(self.NAMESPACED_FLAG_NAME, active=data['waffle_enabled']):
+                # check twice to test that the result is properly cached
+                self.assertEqual(self.TEST_ORG_FLAG.is_enabled(self.TEST_ORG_KEY), data['result'])
+                self.assertEqual(self.TEST_ORG_FLAG.is_enabled(self.TEST_ORG_KEY), data['result'])
+                # result is cached, so override check should happen once
+                WaffleFlagOrgOverrideModel.override_value.assert_called_once_with(
+                    self.NAMESPACED_FLAG_NAME,
+                    self.TEST_ORG_KEY
+                )
+
+        # check flag for a second org
+        if data['org_override'] == WaffleFlagOrgOverrideModel.ALL_CHOICES.unset:
+            # When org override wasn't set for the first org, the second org will get the same
+            # cached value from waffle.
+            self.assertEqual(self.TEST_ORG_FLAG.is_enabled(self.TEST_ORG_2_KEY), data['waffle_enabled'])
+        else:
+            # When org override was set for the first org, it should not apply to the second
+            # org which should get the default value of False.
+            self.assertEqual(self.TEST_ORG_FLAG.is_enabled(self.TEST_ORG_2_KEY), False)
+
+    @ddt.data(
+        {'flag_undefined_default': None, 'result': False},
+        {'flag_undefined_default': False, 'result': False},
+        {'flag_undefined_default': True, 'result': True},
+    )
+    def test_undefined_waffle_flag(self, data):
+        """
+        Test flag with various defaults provided for undefined waffle flags.
+        """
+        RequestCache.clear_request_cache()
+
+        test_org_flag = OrgWaffleFlag(
+            self.TEST_NAMESPACE,
+            self.FLAG_NAME,
+            flag_undefined_default=data['flag_undefined_default']
+        )
+
+        with patch.object(
+            WaffleFlagOrgOverrideModel,
+            'override_value',
+            return_value=WaffleFlagOrgOverrideModel.ALL_CHOICES.unset
+        ):
+            # check twice to test that the result is properly cached
+            self.assertEqual(test_org_flag.is_enabled(self.TEST_ORG_KEY), data['result'])
+            self.assertEqual(test_org_flag.is_enabled(self.TEST_ORG_KEY), data['result'])
+            # result is cached, so override check should happen once
+            WaffleFlagOrgOverrideModel.override_value.assert_called_once_with(
+                self.NAMESPACED_FLAG_NAME,
+                self.TEST_ORG_KEY
+            )

openedx/core/djangoapps/waffle_utils/tests/test_models.py

@@ -7,7 +7,7 @@ from opaque_keys.edx.keys import CourseKey
 
 from request_cache.middleware import RequestCache
 
-from ..models import WaffleFlagCourseOverrideModel
+from ..models import WaffleFlagCourseOverrideModel, WaffleFlagOrgOverrideModel
 
 
 @ddt
@@ -49,3 +49,44 @@ class WaffleFlagCourseOverrideTests(TestCase):
             enabled=is_enabled,
             course_id=self.TEST_COURSE_KEY
         )
+
+
+@ddt
+class WaffleFlagOrgOverrideTests(TestCase):
+    """
+    Tests for the waffle flag organization override model.
+    """
+
+    WAFFLE_TEST_NAME = 'waffle_test_org_override'
+    TEST_ORG_KEY = 'edX'
+    OVERRIDE_CHOICES = WaffleFlagOrgOverrideModel.ALL_CHOICES
+
+    # Data format: ( is_enabled, override_choice, expected_result )
+    @data((True, OVERRIDE_CHOICES.on, OVERRIDE_CHOICES.on),
+          (True, OVERRIDE_CHOICES.off, OVERRIDE_CHOICES.off),
+          (False, OVERRIDE_CHOICES.on, OVERRIDE_CHOICES.unset))
+    @unpack
+    def test_setting_override(self, is_enabled, override_choice, expected_result):
+        RequestCache.clear_request_cache()
+        self.set_waffle_org_override(override_choice, is_enabled)
+        override_value = WaffleFlagOrgOverrideModel.override_value(
+            self.WAFFLE_TEST_NAME, self.TEST_ORG_KEY
+        )
+        self.assertEqual(override_value, expected_result)
+
+    def test_setting_override_multiple_times(self):
+        RequestCache.clear_request_cache()
+        self.set_waffle_org_override(self.OVERRIDE_CHOICES.on)
+        self.set_waffle_org_override(self.OVERRIDE_CHOICES.off)
+        override_value = WaffleFlagOrgOverrideModel.override_value(
+            self.WAFFLE_TEST_NAME, self.TEST_ORG_KEY
+        )
+        self.assertEqual(override_value, self.OVERRIDE_CHOICES.off)
+
+    def set_waffle_org_override(self, override_choice, is_enabled=True):
+        WaffleFlagOrgOverrideModel.objects.create(
+            waffle_flag=self.WAFFLE_TEST_NAME,
+            override_choice=override_choice,
+            enabled=is_enabled,
+            org_id=self.TEST_ORG_KEY
+        )

openedx/core/djangoapps/waffle_utils/testutils.py

@@ -10,7 +10,8 @@ from waffle.testutils import override_flag
 # waffle tables. For example:
 #   QUERY_COUNT_TABLE_BLACKLIST = WAFFLE_TABLES
 #   with self.assertNumQueries(6, table_blacklist=QUERY_COUNT_TABLE_BLACKLIST):
-WAFFLE_TABLES = ['waffle_utils_waffleflagcourseoverridemodel', 'waffle_flag', 'waffle_switch', 'waffle_sample']
+WAFFLE_TABLES = ['waffle_utils_waffleflagcourseoverridemodel', 'waffle_utils_waffleflagorgoverridemodel', 'waffle_flag',
+                 'waffle_switch', 'waffle_sample']
 
 
 def override_waffle_flag(flag, active):