Convert AdminFlagValue to enum

tests/unit/accounts/test_views.py

@@ -904,7 +904,7 @@ def test_register_redirect(self, db_request, monkeypatch):
     def test_register_fails_with_admin_flag_set(self, db_request):
         # This flag was already set via migration, just need to enable it
         flag = db_request.db.query(AdminFlag).get(
-            AdminFlagValue.DISALLOW_NEW_USER_REGISTRATION
+            AdminFlagValue.DISALLOW_NEW_USER_REGISTRATION.value
         )
         flag.enabled = True
 

tests/unit/admin/test_flags.py

@@ -10,14 +10,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import enum
+
 from ...common.db.admin import AdminFlagFactory
 
 
+class TestAdminFlagValues(enum.Enum):
+    NOT_A_REAL_FLAG = "not-a-real-flag"
+    THIS_FLAG_IS_ENABLED = "this-flag-is-enabled"
+
+
 class TestAdminFlag:
     def test_default(self, db_request):
-        assert not db_request.flags.enabled("not-a-real-flag")
+        assert not db_request.flags.enabled(TestAdminFlagValues.NOT_A_REAL_FLAG)
 
     def test_enabled(self, db_request):
         AdminFlagFactory(id="this-flag-is-enabled")
 
-        assert db_request.flags.enabled("this-flag-is-enabled")
+        assert db_request.flags.enabled(TestAdminFlagValues.THIS_FLAG_IS_ENABLED)

tests/unit/forklift/test_legacy.py

@@ -1137,7 +1137,9 @@ def test_fails_with_stdlib_names(self, pyramid_config, db_request, name):
     def test_fails_with_admin_flag_set(self, pyramid_config, db_request):
         admin_flag = (
             db_request.db.query(AdminFlag)
-            .filter(AdminFlag.id == AdminFlagValue.DISALLOW_NEW_PROJECT_REGISTRATION)
+            .filter(
+                AdminFlag.id == AdminFlagValue.DISALLOW_NEW_PROJECT_REGISTRATION.value
+            )
             .first()
         )
         admin_flag.enabled = True

tests/unit/test_db.py

@@ -274,7 +274,7 @@ def test_create_session_read_only_mode(
     )
 
     assert _create_session(request) is session_obj
-    assert get.calls == [pretend.call(AdminFlagValue.READ_ONLY)]
+    assert get.calls == [pretend.call(AdminFlagValue.READ_ONLY.value)]
     assert request.tm.doom.calls == doom_calls
 
 

warehouse/admin/flags.py

@@ -10,12 +10,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import enum
+
 from sqlalchemy import Boolean, Column, Text, sql
 
 from warehouse import db
 
 
-class AdminFlagValue:
+class AdminFlagValue(enum.Enum):
     DISALLOW_DELETION = "disallow-deletion"
     DISALLOW_NEW_PROJECT_REGISTRATION = "disallow-new-project-registration"
     DISALLOW_NEW_UPLOAD = "disallow-new-upload"
@@ -44,8 +46,8 @@ def notifications(self):
             .all()
         )
 
-    def enabled(self, flag_name):
-        flag = self.request.db.query(AdminFlag).get(flag_name)
+    def enabled(self, flag_member):
+        flag = self.request.db.query(AdminFlag).get(flag_member.value)
         return flag.enabled if flag else False
 
 

warehouse/db.py

@@ -210,7 +210,7 @@ def cleanup(request):
     # Check if we're in read-only mode
     from warehouse.admin.flags import AdminFlag, AdminFlagValue
 
-    flag = session.query(AdminFlag).get(AdminFlagValue.READ_ONLY)
+    flag = session.query(AdminFlag).get(AdminFlagValue.READ_ONLY.value)
     if flag and flag.enabled and not request.user.is_superuser:
         request.tm.doom()