From 556930524886b20478a52a04a02bfa50a18b30e9 Mon Sep 17 00:00:00 2001 From: Sunny Aggarwal Date: Wed, 21 Oct 2026 09:44:18 +0530 Subject: [PATCH 1/2] OpenConceptLab/ocl_online#517 | capacity tier per group from the groups config, output language permission --- .../commands/seed_group_capabilities.py | 23 +++++++++++-- core/capabilities/tests/groups.test.yaml | 2 ++ core/capabilities/tests/tests.py | 17 ++++++++++ core/capacity/constants.py | 4 ++- core/capacity/limiter.py | 19 ++++++----- .../migrations/0002_groupcapacitytier.py | 24 ++++++++++++++ core/capacity/models.py | 13 ++++++++ core/capacity/tests/tests.py | 32 +++++++++++++++++-- core/fixtures/auth_groups.yaml | 6 ++++ 9 files changed, 124 insertions(+), 16 deletions(-) create mode 100644 core/capacity/migrations/0002_groupcapacitytier.py diff --git a/core/capabilities/management/commands/seed_group_capabilities.py b/core/capabilities/management/commands/seed_group_capabilities.py index 9e923e445..b6e071bb2 100644 --- a/core/capabilities/management/commands/seed_group_capabilities.py +++ b/core/capabilities/management/commands/seed_group_capabilities.py @@ -5,12 +5,17 @@ from django.core.management import BaseCommand, CommandError from django.db import transaction +from core.capacity.constants import TIERS + GROUPS_CONFIG_ENV = 'OCL_GROUPS_CONFIG' -# Plan groups are deployment config: {"groups": {"": {"permissions": [...], "capabilities": {...}}}} +# Groups are deployment config: +# {"groups": {"": {"permissions": [...], "capabilities": {...}, "capacity_tier": ""}}} # Sync - the config is the source of truth for every group it declares: users.* permissions are set exactly, -# declared limits are upserted and undeclared ones deleted. Groups it doesn't declare are left alone. +# declared limits are upserted and undeclared ones deleted, and so is the capacity tier. Groups it doesn't declare +# are left alone. # Limits: 0 = unlimited, -1 = blocked (the kill switch), N > 0 = the limit; no row = not entitled (blocked). +# capacity_tier: one of core.capacity.constants.TIERS; a user gets the highest of their groups' (core.capacity.limiter). def seed_groups(config, stdout=None): @@ -47,6 +52,17 @@ def seed_group(group_name, group_config, warn): GroupCapability.objects.update_or_create(group=group, capability=capability, defaults={'limit': limit}) declared_ids.append(capability.id) GroupCapability.objects.filter(group=group).exclude(capability_id__in=declared_ids).delete() + seed_capacity_tier(group, group_config.get('capacity_tier')) + + +def seed_capacity_tier(group, tier): + from core.capacity.models import GroupCapacityTier + + # sync: the declared capacity tier is upserted, or the group's is deleted + if tier: + GroupCapacityTier.objects.update_or_create(group=group, defaults={'tier': tier}) + else: + GroupCapacityTier.objects.filter(group=group).delete() def validate_config(config): @@ -66,6 +82,9 @@ def validate_config(config): not all(isinstance(limit, int) and not isinstance(limit, bool) for limit in capabilities.values()) ): raise CommandError(f'{group_name}: capabilities must map names to integer limits') + tier = group_config.get('capacity_tier') + if tier is not None and tier not in TIERS: + raise CommandError(f'{group_name}: capacity_tier must be one of {", ".join(TIERS)}') return config['groups'] diff --git a/core/capabilities/tests/groups.test.yaml b/core/capabilities/tests/groups.test.yaml index c676c6763..1d588c72c 100644 --- a/core/capabilities/tests/groups.test.yaml +++ b/core/capabilities/tests/groups.test.yaml @@ -1,6 +1,7 @@ groups: preview: permissions: [mapper_use, mapper_ai_assistant] + capacity_tier: preview capabilities: mapper.projects: 1 mapper.rows_per_project: 25 @@ -16,3 +17,4 @@ groups: clone.resources_per_call: 1000 early_access: permissions: [mapper_use, mapper_ai_assistant, mapper_scispacy] + capacity_tier: early_access diff --git a/core/capabilities/tests/tests.py b/core/capabilities/tests/tests.py index fa84470fd..43e25b497 100644 --- a/core/capabilities/tests/tests.py +++ b/core/capabilities/tests/tests.py @@ -410,6 +410,23 @@ def test_creates_group_with_permissions_and_limits(self): self.assertEqual(self.codenames('pro'), {'mapper_use', 'mapper_custom_algorithms', 'mapper_scispacy'}) self.assertEqual(self.limits('pro'), {'mapper.projects': 20, 'mapper.match_operations': 0}) + def test_seeds_and_syncs_the_capacity_tier(self): + from core.capacity.models import GroupCapacityTier + config = {'groups': {'group_a': {'capacity_tier': 'early_access'}, 'group_b': {}}} + self.seed(config) + self.assertEqual(GroupCapacityTier.objects.get(group__name='group_a').tier, 'early_access') + self.assertFalse(GroupCapacityTier.objects.filter(group__name='group_b').exists()) + + self.seed({'groups': {'group_a': {'capacity_tier': 'core'}}}) + self.assertEqual(GroupCapacityTier.objects.get(group__name='group_a').tier, 'core') + + self.seed({'groups': {'group_a': {'permissions': ['mapper_use']}}}) + self.assertFalse(GroupCapacityTier.objects.filter(group__name='group_a').exists()) + + def test_rejects_an_unknown_capacity_tier(self): + with self.assertRaisesMessage(CommandError, 'group_a: capacity_tier must be one of'): + self.seed({'groups': {'group_a': {'capacity_tier': 'gold'}}}) + def test_reseed_syncs_permissions_and_limits(self): self.seed(self.CONFIG) self.seed({ diff --git a/core/capacity/constants.py b/core/capacity/constants.py index 248ae015f..4c9050732 100644 --- a/core/capacity/constants.py +++ b/core/capacity/constants.py @@ -9,12 +9,14 @@ ENFORCE_FOR = (ENFORCE_FOR_AWARE, ENFORCE_FOR_ALL) CAPACITY_AWARE_METADATA_KEY = 'capacity_aware' -# A user's tier is their highest plan group, highest first. +# Capacity tiers, highest first. Staff and superusers get the staff tier; everyone else the highest tier their groups +# are set to (`capacity_tier` in the groups config, core.capacity.models.GroupCapacityTier), else the lowest. TIER_STAFF = 'staff' TIER_CORE = 'core' TIER_EARLY_ACCESS = 'early_access' TIER_PREVIEW = 'preview' TIERS = (TIER_STAFF, TIER_CORE, TIER_EARLY_ACCESS, TIER_PREVIEW) +DEFAULT_TIER = TIERS[-1] # Lanes: each counts the heavy calls in flight in one scope. LANE_API_HEAVY = 'api_heavy' # every heavy call, cluster-wide diff --git a/core/capacity/limiter.py b/core/capacity/limiter.py index 5c3426255..2bb9d7512 100644 --- a/core/capacity/limiter.py +++ b/core/capacity/limiter.py @@ -34,15 +34,14 @@ from core.capacity.config import get_config from core.capacity.constants import ( - TIERS, MODE_OFF, MODE_ENFORCE, ENFORCE_FOR_ALL, CAPACITY_AWARE_METADATA_KEY, TIER_STAFF, TIER_CORE, - TIER_EARLY_ACCESS, TIER_PREVIEW, LANE_API_HEAVY, LANE_API_HEAVY_TASK, LANE_ES2_KNN, LANE_TIER, LANE_USER, + TIERS, DEFAULT_TIER, MODE_OFF, MODE_ENFORCE, ENFORCE_FOR_ALL, CAPACITY_AWARE_METADATA_KEY, TIER_STAFF, + LANE_API_HEAVY, LANE_API_HEAVY_TASK, LANE_ES2_KNN, LANE_TIER, LANE_USER, DECISION_ADMITTED, DECISION_SHADOW_REFUSED, DECISION_REFUSED, DECISION_UNAVAILABLE, ENDPOINT_MATCH, HEADER_DECISION, HEADER_LIMIT, HEADER_IN_FLIGHT, HEADER_TIER, HEADER_TIER_LIMIT, HEADER_TIER_IN_FLIGHT, HEADER_SUGGESTED_CONCURRENCY, CAPACITY_EXCEEDED_ERROR_CODE, LOG_EVENT, REDIS_KEY_PREFIX) from core.capacity.logs import emit from core.capacity.threads import RedisThread from core.common.utils import get_event_metadata -from core.users.constants import CORE_USER_GROUP, EARLY_ACCESS_GROUP _NOW_MS = """ if redis.replicate_commands then redis.replicate_commands() end @@ -248,15 +247,15 @@ def get_not_after_ms(): def get_tier(user): - """The user's highest plan tier, as capabilities resolve it: staff > core > early_access > preview.""" + """ + The user's capacity tier: staff for staff and superusers, else the highest tier their groups are set to + (`capacity_tier` in the groups config, OpenConceptLab/ocl_online#517), else the lowest. + """ + from core.capacity.models import GroupCapacityTier if user.is_staff or user.is_superuser: return TIER_STAFF - groups = set(user.groups.values_list('name', flat=True)) - if CORE_USER_GROUP in groups: - return TIER_CORE - if EARLY_ACCESS_GROUP in groups: - return TIER_EARLY_ACCESS - return TIER_PREVIEW + tiers = set(GroupCapacityTier.objects.filter(group__in=user.groups.all()).values_list('tier', flat=True)) + return next((tier for tier in TIERS if tier in tiers), DEFAULT_TIER) def get_queue_ms(request, now=None): diff --git a/core/capacity/migrations/0002_groupcapacitytier.py b/core/capacity/migrations/0002_groupcapacitytier.py new file mode 100644 index 000000000..ad581a2ab --- /dev/null +++ b/core/capacity/migrations/0002_groupcapacitytier.py @@ -0,0 +1,24 @@ +import django.db.models.deletion +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('auth', '0012_alter_user_first_name_max_length'), + ('capacity', '0001_initial'), + ] + + operations = [ + migrations.CreateModel( + name='GroupCapacityTier', + fields=[ + ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('tier', models.CharField(max_length=32)), + ('group', models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, related_name='capacity_tier', to='auth.group')), + ], + options={ + 'db_table': 'group_capacity_tiers', + }, + ), + ] diff --git a/core/capacity/models.py b/core/capacity/models.py index 9ec566485..369508a23 100644 --- a/core/capacity/models.py +++ b/core/capacity/models.py @@ -20,3 +20,16 @@ class Meta: @classmethod def get_latest(cls): return cls.objects.order_by('-id').first() + + +class GroupCapacityTier(models.Model): + """ + The capacity tier a group's members get (OpenConceptLab/ocl_online#517), from `capacity_tier` in the groups + config, which `seed_group_capabilities` syncs. A user gets the highest tier of their groups; core.capacity.limiter + never reads group names. + """ + class Meta: + db_table = 'group_capacity_tiers' + + group = models.OneToOneField('auth.Group', on_delete=models.CASCADE, related_name='capacity_tier') + tier = models.CharField(max_length=32) diff --git a/core/capacity/tests/tests.py b/core/capacity/tests/tests.py index 0e1699b15..c2397e23a 100644 --- a/core/capacity/tests/tests.py +++ b/core/capacity/tests/tests.py @@ -28,7 +28,7 @@ from core.capacity.limiter import ( RedisLanes, CapacityGate, LeaseRenewer, get_tier, get_queue_ms, lane_key, get_task_id, build_redis_client, get_status) -from core.capacity.models import CapacityConfig +from core.capacity.models import CapacityConfig, GroupCapacityTier from core.capacity.threads import RedisThread from core.common.tests import OCLTestCase, OCLAPITestCase, PREVIEW_GROUP_NAME from core.common.utils import get_event_metadata @@ -36,10 +36,22 @@ from core.users.tests.factories import UserProfileFactory +# The tiers the groups config sets for these test groups (seed_group_capabilities syncs them in deployments). +TEST_GROUP_TIERS = {CORE_USER_GROUP: 'core', EARLY_ACCESS_GROUP: 'early_access', PREVIEW_GROUP_NAME: 'preview'} + + +def get_group(name, tier=None): + group = Group.objects.get_or_create(name=name)[0] + tier = tier or TEST_GROUP_TIERS.get(name) + if tier: + GroupCapacityTier.objects.update_or_create(group=group, defaults={'tier': tier}) + return group + + def make_user(group=None, **kwargs): user = UserProfileFactory(**kwargs) if group: - user.groups.add(Group.objects.get_or_create(name=group)[0]) + user.groups.add(get_group(group)) return user @@ -242,7 +254,21 @@ def test_get_tier(self): self.assertEqual(get_tier(make_user(PREVIEW_GROUP_NAME)), 'preview') self.assertEqual(get_tier(make_user()), 'preview') user = make_user(PREVIEW_GROUP_NAME) - user.groups.add(Group.objects.get(name=CORE_USER_GROUP)) + user.groups.add(get_group(CORE_USER_GROUP)) + self.assertEqual(get_tier(user), 'core') + + def test_get_tier_comes_from_the_groups_tier_not_its_name(self): + get_group('group_a', tier='early_access') + self.assertEqual(get_tier(make_user('group_a')), 'early_access') + self.assertEqual(get_tier(make_user('group_without_tier')), 'preview') + user = make_user(CORE_USER_GROUP) + GroupCapacityTier.objects.filter(group__name=CORE_USER_GROUP).delete() + self.assertEqual(get_tier(user), 'preview') + + def test_get_tier_is_the_highest_of_the_users_groups(self): + user = make_user('group_a') + user.groups.add(get_group('group_b', tier='core')) + user.groups.add(get_group('group_c', tier='preview')) self.assertEqual(get_tier(user), 'core') def test_get_queue_ms(self): diff --git a/core/fixtures/auth_groups.yaml b/core/fixtures/auth_groups.yaml index 7b9c0a3d7..38a61b4e6 100644 --- a/core/fixtures/auth_groups.yaml +++ b/core/fixtures/auth_groups.yaml @@ -54,6 +54,12 @@ codename: list_unpaginated name: Can request unpaginated (Compress) list responses content_type: [users, userprofile] +- model: "auth.permission" + pk: 1009 + fields: + codename: mapper_ai_output_locale + name: Can choose the Mapper AI Assistant's output language + content_type: [users, userprofile] - model: "auth.group" pk: 1 fields: From f850957681d725d818d419991e5b90faa66dcb24 Mon Sep 17 00:00:00 2001 From: Sunny Aggarwal Date: Thu, 22 Oct 2026 07:29:19 +0530 Subject: [PATCH 2/2] OpenConceptLab/ocl_online#517 | Permission fixture by natural key, neutral tier groups in capacity tests --- core/capabilities/tests/groups.test.yaml | 2 - core/capacity/tests/tests.py | 124 +++++++++++------------ core/fixtures/auth_groups.yaml | 3 +- 3 files changed, 64 insertions(+), 65 deletions(-) diff --git a/core/capabilities/tests/groups.test.yaml b/core/capabilities/tests/groups.test.yaml index 1d588c72c..c676c6763 100644 --- a/core/capabilities/tests/groups.test.yaml +++ b/core/capabilities/tests/groups.test.yaml @@ -1,7 +1,6 @@ groups: preview: permissions: [mapper_use, mapper_ai_assistant] - capacity_tier: preview capabilities: mapper.projects: 1 mapper.rows_per_project: 25 @@ -17,4 +16,3 @@ groups: clone.resources_per_call: 1000 early_access: permissions: [mapper_use, mapper_ai_assistant, mapper_scispacy] - capacity_tier: early_access diff --git a/core/capacity/tests/tests.py b/core/capacity/tests/tests.py index c2397e23a..142e8a01e 100644 --- a/core/capacity/tests/tests.py +++ b/core/capacity/tests/tests.py @@ -23,7 +23,7 @@ HEADERS, HEADER_DECISION, HEADER_LIMIT, HEADER_IN_FLIGHT, HEADER_TIER, HEADER_TIER_LIMIT, HEADER_TIER_IN_FLIGHT, HEADER_SUGGESTED_CONCURRENCY, ENDPOINT_MATCH, ENDPOINT_RERANK, DECISION_ADMITTED, DECISION_SHADOW_REFUSED, DECISION_REFUSED, DECISION_UNAVAILABLE, LANE_API_HEAVY, LANE_API_HEAVY_TASK, LANE_ES2_KNN, LANE_TIER, LANE_USER, - LOG_EVENT, CONFIG_LOG_EVENT, SOURCE_COMMAND) + LOG_EVENT, CONFIG_LOG_EVENT, SOURCE_COMMAND, TIER_CORE, TIER_EARLY_ACCESS, TIER_PREVIEW) from core.capacity import logs from core.capacity.limiter import ( RedisLanes, CapacityGate, LeaseRenewer, get_tier, get_queue_ms, lane_key, get_task_id, build_redis_client, @@ -32,26 +32,26 @@ from core.capacity.threads import RedisThread from core.common.tests import OCLTestCase, OCLAPITestCase, PREVIEW_GROUP_NAME from core.common.utils import get_event_metadata -from core.users.constants import CORE_USER_GROUP, EARLY_ACCESS_GROUP from core.users.tests.factories import UserProfileFactory -# The tiers the groups config sets for these test groups (seed_group_capabilities syncs them in deployments). -TEST_GROUP_TIERS = {CORE_USER_GROUP: 'core', EARLY_ACCESS_GROUP: 'early_access', PREVIEW_GROUP_NAME: 'preview'} - - def get_group(name, tier=None): group = Group.objects.get_or_create(name=name)[0] - tier = tier or TEST_GROUP_TIERS.get(name) if tier: GroupCapacityTier.objects.update_or_create(group=group, defaults={'tier': tier}) return group -def make_user(group=None, **kwargs): +def get_tier_group(tier): + return get_group(f'group-{tier}', tier=tier) + + +def make_user(group=None, tier=None, **kwargs): user = UserProfileFactory(**kwargs) if group: user.groups.add(get_group(group)) + if tier: + user.groups.add(get_tier_group(tier)) return user @@ -249,20 +249,20 @@ class CapacityHelpersTest(OCLTestCase): def test_get_tier(self): self.assertEqual(get_tier(make_user(is_staff=True)), 'staff') self.assertEqual(get_tier(make_user(is_superuser=True)), 'staff') - self.assertEqual(get_tier(make_user(CORE_USER_GROUP)), 'core') - self.assertEqual(get_tier(make_user(EARLY_ACCESS_GROUP)), 'early_access') - self.assertEqual(get_tier(make_user(PREVIEW_GROUP_NAME)), 'preview') + self.assertEqual(get_tier(make_user(tier=TIER_CORE)), 'core') + self.assertEqual(get_tier(make_user(tier=TIER_EARLY_ACCESS)), 'early_access') + self.assertEqual(get_tier(make_user(tier=TIER_PREVIEW)), 'preview') self.assertEqual(get_tier(make_user()), 'preview') - user = make_user(PREVIEW_GROUP_NAME) - user.groups.add(get_group(CORE_USER_GROUP)) + user = make_user(tier=TIER_PREVIEW) + user.groups.add(get_tier_group(TIER_CORE)) self.assertEqual(get_tier(user), 'core') def test_get_tier_comes_from_the_groups_tier_not_its_name(self): get_group('group_a', tier='early_access') self.assertEqual(get_tier(make_user('group_a')), 'early_access') self.assertEqual(get_tier(make_user('group_without_tier')), 'preview') - user = make_user(CORE_USER_GROUP) - GroupCapacityTier.objects.filter(group__name=CORE_USER_GROUP).delete() + user = make_user(tier=TIER_CORE) + GroupCapacityTier.objects.filter(group__name='group-core').delete() self.assertEqual(get_tier(user), 'preview') def test_get_tier_is_the_highest_of_the_users_groups(self): @@ -340,7 +340,7 @@ def test_build_redis_client_with_sentinels(self): class CapacityGateTest(CapacityTestMixin, OCLTestCase): def test_admits_and_counts_a_call_in_every_lane(self): - user = make_user(EARLY_ACCESS_GROUP) + user = make_user(tier=TIER_EARLY_ACCESS) gate = self.acquire(user) @@ -362,7 +362,7 @@ def test_admits_and_counts_a_call_in_every_lane(self): self.assertEqual(len(self.lines()), 1) def test_shadow_mode_never_refuses_but_says_it_would_have(self): - user = make_user(PREVIEW_GROUP_NAME) + user = make_user(tier=TIER_PREVIEW) first = self.acquire(user) second = self.acquire(user) @@ -377,7 +377,7 @@ def test_shadow_mode_never_refuses_but_says_it_would_have(self): def test_enforce_mode_refuses_without_taking_a_lease(self): self.configure(mode='enforce', enforce_for='all') - user = make_user(PREVIEW_GROUP_NAME) + user = make_user(tier=TIER_PREVIEW) first = self.acquire(user) second = self.acquire(user) @@ -401,7 +401,7 @@ def test_enforce_mode_refuses_without_taking_a_lease(self): def test_enforce_for_aware_refuses_only_clients_that_say_they_handle_it(self): self.configure(mode='enforce') - user = make_user(PREVIEW_GROUP_NAME) + user = make_user(tier=TIER_PREVIEW) self.acquire(user) unaware = self.acquire(user, metadata={'algorithm_id': 'ocl-semantic'}) @@ -413,7 +413,7 @@ def test_enforce_for_aware_refuses_only_clients_that_say_they_handle_it(self): self.assertTrue(aware.enforced) def test_a_stale_lease_expires(self): - user = make_user(PREVIEW_GROUP_NAME) + user = make_user(tier=TIER_PREVIEW) # a lease left by a worker that died: its expiry has passed self.redis.zadd(lane_key(LANE_USER, user.id), {'dead-worker': 1000}) @@ -424,11 +424,11 @@ def test_a_stale_lease_expires(self): def test_single_row_calls_can_use_the_reserved_slot(self): self.configure(mode='enforce', enforce_for='all', api_heavy={'cluster': 2, 'per_task': 5}) - bulk = self.acquire(make_user(CORE_USER_GROUP)) + bulk = self.acquire(make_user(tier=TIER_CORE)) - second_bulk = self.acquire(make_user(CORE_USER_GROUP)) - single_row = self.acquire(make_user(CORE_USER_GROUP), rows=1) - rerank = self.acquire(make_user(CORE_USER_GROUP), endpoint=ENDPOINT_RERANK, rows=1, semantic=False) + second_bulk = self.acquire(make_user(tier=TIER_CORE)) + single_row = self.acquire(make_user(tier=TIER_CORE), rows=1) + rerank = self.acquire(make_user(tier=TIER_CORE), endpoint=ENDPOINT_RERANK, rows=1, semantic=False) self.assertEqual(bulk.decision, DECISION_ADMITTED) self.assertEqual(bulk.limits[LANE_API_HEAVY], 1) @@ -442,10 +442,10 @@ def test_single_row_calls_can_use_the_reserved_slot(self): def test_each_api_task_has_its_own_lane(self): self.configure(mode='enforce', enforce_for='all', api_heavy={'cluster': 4, 'per_task': 1}, reserve_single_row=0) with override_settings(CAPACITY_TASK_ID='task-a'): - first = self.acquire(make_user(CORE_USER_GROUP)) - second = self.acquire(make_user(CORE_USER_GROUP)) + first = self.acquire(make_user(tier=TIER_CORE)) + second = self.acquire(make_user(tier=TIER_CORE)) with override_settings(CAPACITY_TASK_ID='task-b'): - other_task = self.acquire(make_user(CORE_USER_GROUP)) + other_task = self.acquire(make_user(tier=TIER_CORE)) self.assertEqual(first.decision, DECISION_ADMITTED) self.assertEqual((second.decision, second.full), (DECISION_REFUSED, [LANE_API_HEAVY_TASK])) @@ -454,11 +454,11 @@ def test_each_api_task_has_its_own_lane(self): def test_tier_ceilings_leave_other_tiers_alone(self): self.configure(mode='enforce', enforce_for='all', api_heavy={'cluster': 10, 'per_task': 10}, es2_knn=10) - self.acquire(make_user(PREVIEW_GROUP_NAME)) - self.acquire(make_user(PREVIEW_GROUP_NAME)) + self.acquire(make_user(tier=TIER_PREVIEW)) + self.acquire(make_user(tier=TIER_PREVIEW)) - third_preview = self.acquire(make_user(PREVIEW_GROUP_NAME)) - early_access = self.acquire(make_user(EARLY_ACCESS_GROUP)) + third_preview = self.acquire(make_user(tier=TIER_PREVIEW)) + early_access = self.acquire(make_user(tier=TIER_EARLY_ACCESS)) self.assertEqual((third_preview.decision, third_preview.full), (DECISION_REFUSED, [LANE_TIER])) self.assertEqual(third_preview.counts[LANE_TIER], 2) @@ -467,10 +467,10 @@ def test_tier_ceilings_leave_other_tiers_alone(self): def test_only_semantic_calls_take_an_es2_knn_lane(self): self.configure(mode='enforce', enforce_for='all', es2_knn=1) - semantic = self.acquire(make_user(CORE_USER_GROUP)) + semantic = self.acquire(make_user(tier=TIER_CORE)) - semantic_again = self.acquire(make_user(CORE_USER_GROUP)) - reranked_lexical = self.acquire(make_user(CORE_USER_GROUP), semantic=False, reranker=True) + semantic_again = self.acquire(make_user(tier=TIER_CORE)) + reranked_lexical = self.acquire(make_user(tier=TIER_CORE), semantic=False, reranker=True) self.assertIn(LANE_ES2_KNN, semantic.limits) self.assertEqual((semantic_again.decision, semantic_again.full), (DECISION_REFUSED, [LANE_ES2_KNN])) @@ -480,7 +480,7 @@ def test_only_semantic_calls_take_an_es2_knn_lane(self): def test_a_paused_tier(self): self.configure(mode='enforce', enforce_for='all', tiers={'preview': 0}) - gate = self.acquire(make_user(PREVIEW_GROUP_NAME)) + gate = self.acquire(make_user(tier=TIER_PREVIEW)) self.assertEqual((gate.decision, gate.full), (DECISION_REFUSED, [LANE_TIER])) self.assertEqual(gate.retry_after, 120) @@ -490,9 +490,9 @@ def test_a_paused_tier(self): def test_a_paused_lane_is_the_scope_reported(self): self.configure(mode='enforce', enforce_for='all', reserve_single_row=0, api_heavy={'cluster': 1, 'per_task': 5}, tiers={'preview': 0}) - self.acquire(make_user(CORE_USER_GROUP)) + self.acquire(make_user(tier=TIER_CORE)) - gate = self.acquire(make_user(PREVIEW_GROUP_NAME)) + gate = self.acquire(make_user(tier=TIER_PREVIEW)) self.assertEqual(gate.full, [LANE_API_HEAVY, LANE_TIER]) self.assertEqual(gate.scope, LANE_TIER) @@ -502,7 +502,7 @@ def test_a_paused_lane_is_the_scope_reported(self): def test_retry_after_grows_with_the_calls_ahead_up_to_the_max(self): self.configure(tiers={'core': 10}, per_user={'core': 10}, api_heavy={'cluster': 10, 'per_task': 10}) - user = make_user(CORE_USER_GROUP) + user = make_user(tier=TIER_CORE) gates = [self.acquire(user, semantic=False, reranker=True) for _ in range(12)] # bulk calls get 9 of the cluster's 10 @@ -511,18 +511,18 @@ def test_retry_after_grows_with_the_calls_ahead_up_to_the_max(self): self.assertEqual(self.acquire(user, semantic=False, reranker=True).retry_after, 12) def test_suggested_concurrency(self): - core = make_user(CORE_USER_GROUP) + core = make_user(tier=TIER_CORE) first = self.acquire(core) self.assertEqual(first.get_suggested_concurrency(), 3) # alone: up to the per-user limit - second = self.acquire(make_user(CORE_USER_GROUP)) + second = self.acquire(make_user(tier=TIER_CORE)) self.assertEqual(second.get_suggested_concurrency(), 2) # its own + the 1 bulk slot left in the cluster - preview = self.acquire(make_user(PREVIEW_GROUP_NAME)) + preview = self.acquire(make_user(tier=TIER_PREVIEW)) self.assertEqual(preview.get_suggested_concurrency(), 1) # never below 1 unless paused self.assertIsNone(CapacityGate(Mock(), ENDPOINT_MATCH).get_suggested_concurrency()) def test_renewal_extends_every_lease_or_none(self): - gate = self.acquire(make_user(PREVIEW_GROUP_NAME)) + gate = self.acquire(make_user(tier=TIER_PREVIEW)) key = lane_key(LANE_USER, gate.request.user.id) soon = time.time() * 1000 + 5000 for lane_key_ in gate.keys: @@ -578,7 +578,7 @@ def test_the_renewer_thread_renews_until_stopped(self): gate.renew.assert_not_called() def test_a_call_starts_a_renewer_and_release_stops_it(self): - gate = self.acquire(make_user(PREVIEW_GROUP_NAME)) + gate = self.acquire(make_user(tier=TIER_PREVIEW)) self.assertTrue(gate.renewer.is_alive()) gate.release() @@ -589,24 +589,24 @@ def test_redis_failure_fails_open_and_skips_redis_for_a_while(self): broken.register_script.return_value = Mock(side_effect=redis.ConnectionError('refused')) RedisLanes.use_client(broken) - gate = self.acquire(make_user(PREVIEW_GROUP_NAME)) + gate = self.acquire(make_user(tier=TIER_PREVIEW)) self.assertEqual(gate.decision, DECISION_UNAVAILABLE) self.assertFalse(gate.holding) self.assertEqual(gate.error, 'ConnectionError: refused') self.assertFalse(RedisLanes.is_available()) - skipped = self.acquire(make_user(PREVIEW_GROUP_NAME)) + skipped = self.acquire(make_user(tier=TIER_PREVIEW)) self.assertEqual(skipped.decision, DECISION_UNAVAILABLE) self.assertEqual(broken.register_script.return_value.call_count, 1) self.assertEqual(skipped.get_suggested_concurrency(), 1) RedisLanes.unavailable_until = 0.0 - self.acquire(make_user(PREVIEW_GROUP_NAME)) + self.acquire(make_user(tier=TIER_PREVIEW)) self.assertEqual(broken.register_script.return_value.call_count, 2) def test_a_failed_release_or_renewal_is_logged_not_raised(self): - gate = self.acquire(make_user(PREVIEW_GROUP_NAME)) + gate = self.acquire(make_user(tier=TIER_PREVIEW)) with patch.object(RedisLanes, 'renew', side_effect=redis.TimeoutError('slow')): gate.renew() self.assertEqual(gate.error, 'TimeoutError: slow') @@ -619,7 +619,7 @@ def test_a_failed_release_or_renewal_is_logged_not_raised(self): self.assertFalse(RedisLanes.is_available()) def test_renewal_and_release_skip_redis_after_it_failed(self): - gate = self.acquire(make_user(PREVIEW_GROUP_NAME)) + gate = self.acquire(make_user(tier=TIER_PREVIEW)) RedisLanes.mark_unavailable() with patch.object(RedisLanes, 'renew') as renew_mock, patch.object(RedisLanes, 'release') as release_mock: @@ -633,7 +633,7 @@ def test_renewal_and_release_skip_redis_after_it_failed(self): def test_a_renewer_that_fails_to_start_still_releases(self): with patch.object(LeaseRenewer, 'start', side_effect=RuntimeError("can't start new thread")): - gate = self.acquire(make_user(PREVIEW_GROUP_NAME)) + gate = self.acquire(make_user(tier=TIER_PREVIEW)) self.assertEqual(gate.decision, DECISION_ADMITTED) self.assertIsNone(gate.renewer) @@ -651,7 +651,7 @@ def test_an_acquire_that_redis_runs_too_late_takes_no_lease(self): self.assertEqual(self.redis.zcard(key), 0) with patch('core.capacity.limiter.get_not_after_ms', return_value=int(time.time() * 1000) - 5000): - gate = self.acquire(make_user(PREVIEW_GROUP_NAME)) + gate = self.acquire(make_user(tier=TIER_PREVIEW)) self.assertEqual(gate.decision, DECISION_UNAVAILABLE) self.assertFalse(gate.holding) self.assertIn('no lease was taken', gate.error) @@ -661,7 +661,7 @@ def test_a_short_lease_never_shortens_a_lanes_ttl(self): self.redis.zadd(lane_key(LANE_API_HEAVY), {'long-lease': time.time() * 1000 + 600000}) self.redis.pexpire(lane_key(LANE_API_HEAVY), 1200000) - gate = self.acquire(make_user(PREVIEW_GROUP_NAME)) + gate = self.acquire(make_user(tier=TIER_PREVIEW)) gate.renew() self.assertGreater(self.redis.pttl(lane_key(LANE_API_HEAVY)), 600000) @@ -674,7 +674,7 @@ def hang(*_): with patch.object(RedisLanes, 'run_script', side_effect=hang): started = time.monotonic() - gate = self.acquire(make_user(PREVIEW_GROUP_NAME)) + gate = self.acquire(make_user(tier=TIER_PREVIEW)) elapsed = time.monotonic() - started self.assertLess(elapsed, 0.4) @@ -703,7 +703,7 @@ def test_an_error_after_a_refusal_never_refuses(self): self.configure(mode='enforce', enforce_for='all', per_user={'preview': 0}) with patch.object(CapacityGate, 'get_retry_after', side_effect=ValueError('bad math')): - gate = self.acquire(make_user(PREVIEW_GROUP_NAME)) + gate = self.acquire(make_user(tier=TIER_PREVIEW)) self.assertEqual(gate.decision, DECISION_UNAVAILABLE) self.assertFalse(gate.refused) @@ -711,7 +711,7 @@ def test_an_error_after_a_refusal_never_refuses(self): def test_an_unexpected_error_fails_open(self): with patch('core.capacity.limiter.get_tier', side_effect=KeyError('tier')): - gate = self.acquire(make_user(PREVIEW_GROUP_NAME)) + gate = self.acquire(make_user(tier=TIER_PREVIEW)) self.assertEqual(gate.decision, DECISION_UNAVAILABLE) self.assertEqual(gate.error, "KeyError: 'tier'") @@ -722,7 +722,7 @@ def test_an_unexpected_error_fails_open(self): def test_off_mode_does_nothing(self): self.configure(mode='off') with patch.object(RedisLanes, 'acquire') as acquire_mock: - gate = self.acquire(make_user(PREVIEW_GROUP_NAME)) + gate = self.acquire(make_user(tier=TIER_PREVIEW)) response = {} gate.apply_headers(response) @@ -735,7 +735,7 @@ def test_off_mode_does_nothing(self): def test_headers(self): self.configure(mode='enforce', enforce_for='all') - user = make_user(PREVIEW_GROUP_NAME) + user = make_user(tier=TIER_PREVIEW) admitted, refused = self.acquire(user), self.acquire(user) admitted_headers, refused_headers = {}, {} @@ -755,7 +755,7 @@ def test_headers(self): self.assertEqual(admitted.error, 'ValueError: bad header') def test_the_log_line(self): - user = make_user(PREVIEW_GROUP_NAME) + user = make_user(tier=TIER_PREVIEW) self.acquire(user) received = int(time.time()) - 3 gate = self.acquire( @@ -780,7 +780,7 @@ def test_the_log_line(self): def test_a_failing_log_line_never_fails_the_call(self): self.limiter_emit.side_effect = OSError('stdout closed') - gate = self.acquire(make_user(PREVIEW_GROUP_NAME)) + gate = self.acquire(make_user(tier=TIER_PREVIEW)) gate.release() @@ -788,8 +788,8 @@ def test_a_failing_log_line_never_fails_the_call(self): def test_status(self): with override_settings(CAPACITY_TASK_ID='task-a'): - gate = self.acquire(make_user(PREVIEW_GROUP_NAME)) - self.acquire(make_user(CORE_USER_GROUP), semantic=False, reranker=True) + gate = self.acquire(make_user(tier=TIER_PREVIEW)) + self.acquire(make_user(tier=TIER_CORE), semantic=False, reranker=True) state = get_status() @@ -809,7 +809,7 @@ class CapacityViewsTest(CapacityTestMixin, OCLAPITestCase): def setUp(self): super().setUp() self.staff = make_user(is_staff=True) - self.preview_user = make_user(PREVIEW_GROUP_NAME) + self.preview_user = make_user(PREVIEW_GROUP_NAME, tier=TIER_PREVIEW) # the group grants mapper_use def post(self, path, data, user=None, **extra): return self.client.post( @@ -1018,7 +1018,7 @@ def test_show_set_reset_and_history(self): self.assertIn(user.username, history) def test_status(self): - self.acquire(make_user(PREVIEW_GROUP_NAME)) + self.acquire(make_user(tier=TIER_PREVIEW)) self.assertEqual(json.loads(self.run_command('status'))[LANE_API_HEAVY]['in_flight'], 1) diff --git a/core/fixtures/auth_groups.yaml b/core/fixtures/auth_groups.yaml index 38a61b4e6..3bd648f1f 100644 --- a/core/fixtures/auth_groups.yaml +++ b/core/fixtures/auth_groups.yaml @@ -54,8 +54,9 @@ codename: list_unpaginated name: Can request unpaginated (Compress) list responses content_type: [users, userprofile] +# No pk: ids from 1009 on belong to permissions Django created for later models, so loaddata +# matches this one by its natural key (codename + content type). - model: "auth.permission" - pk: 1009 fields: codename: mapper_ai_output_locale name: Can choose the Mapper AI Assistant's output language