Allow ".edu.cn" accounts to have more time and disk quota (#4440)

* add edu cn accounts

* add function to check edu account

* add regex to recongnize emails

* fmt

* add unit test

* typo

* add ac.xx to regexes

* fix unittest

codalab/model/bundle_model.py

@@ -68,6 +68,7 @@
 
 SEARCH_KEYWORD_REGEX = re.compile('^([\.\w/]*)=(.*)$')
 SEARCH_RESULTS_LIMIT = 10
+EDU_USER_REGEXES = re.compile('@[\w\.-]+\.(edu|edu\.[a-z]{2}|ac\.[a-z]{2})$')
 
 
 def str_key_dict(row):
@@ -80,6 +81,15 @@ def str_key_dict(row):
     return dict((str(k), v) for k, v in row.items())
 
 
+def is_academic_email(email):
+    """
+    This is a basic function that can be used to compare the email domain suffix with a list of academic email domains.
+    Academic emails typically have domains such as "yy.edu" or "xyz.edu.xx" (where "edu" is followed by a country code).
+    """
+    email_suffix = EDU_USER_REGEXES.findall(email.lower())
+    return len(email_suffix) > 0
+
+
 @dataclass
 class Join:
     """
@@ -2475,18 +2485,19 @@ def add_user(
         :param affiliation:
         :return: (new integer user ID, verification key to send)
         """
+
         with self.engine.begin() as connection:
             now = datetime.datetime.utcnow()
             user_id = user_id or '0x%s' % uuid4().hex
 
             time_quota = (
                 self.default_user_info['edu_time_quota']
-                if email.endswith(".edu")
+                if is_academic_email(email)
                 else self.default_user_info['time_quota']
             )
             disk_quota = (
                 self.default_user_info['edu_disk_quota']
-                if email.endswith(".edu")
+                if is_academic_email(email)
                 else self.default_user_info['disk_quota']
             )
 

tests/unit/model/bundle_model_test.py

@@ -1,6 +1,7 @@
 import unittest
 from tests.unit.server.bundle_manager import TestBase
 from codalab.worker.bundle_state import State
+from codalab.model.bundle_model import is_academic_email
 
 
 class BundleModelTest(TestBase, unittest.TestCase):
@@ -22,6 +23,21 @@ def test_finalizing_bundle_should_not_transition_worker_offline(self):
         bundle = self.bundle_manager._model.get_bundle(bundle.uuid)
         self.assertEqual(bundle.state, State.WORKER_OFFLINE)
 
+    def test_is_academic_email(self):
+        """Unit test to check is_academic_email function."""
+        test_cases = {
+            "[email protected]": True,
+            "[email protected]": True,
+            "[email protected]": True,
+            "[email protected]": False,
+            "[email protected]": False,
+            "[email protected]": True,
+            "[email protected]": True,
+            "[email protected]": False,
+        }
+        for key, value in test_cases.items():
+            self.assertEqual(is_academic_email(key), value)
+
 
 def metadata_to_dicts(uuid, metadata):
     return [