diff --git a/.secrets.baseline b/.secrets.baseline index 95e3cb1a..d7ab771b 100644 --- a/.secrets.baseline +++ b/.secrets.baseline @@ -208,7 +208,7 @@ "filename": "django_app/tests/conftest.py", "hashed_secret": "9a7f654c8f0f9fedc19cf3f0689b661fdb31ee07", "is_verified": false, - "line_number": 156, + "line_number": 166, "is_secret": false } ], @@ -281,5 +281,5 @@ } ] }, - "generated_at": "2026-04-09T11:20:11Z" + "generated_at": "2026-05-05T14:19:26Z" } diff --git a/django_app/redbox_app/redbox_core/migrations/0112_customflag.py b/django_app/redbox_app/redbox_core/migrations/0112_customflag.py index aef364ca..4d34bdf1 100644 --- a/django_app/redbox_app/redbox_core/migrations/0112_customflag.py +++ b/django_app/redbox_app/redbox_core/migrations/0112_customflag.py @@ -9,6 +9,7 @@ class Migration(migrations.Migration): dependencies = [ ("auth", "0012_alter_user_first_name_max_length"), + ("waffle", "0004_update_everyone_nullbooleanfield"), ("redbox_core", "0111_alter_user_ai_experience"), ] diff --git a/django_app/redbox_app/redbox_core/migrations/0113_migrate_waffle_flags.py b/django_app/redbox_app/redbox_core/migrations/0113_migrate_waffle_flags.py new file mode 100644 index 00000000..589cf6b7 --- /dev/null +++ b/django_app/redbox_app/redbox_core/migrations/0113_migrate_waffle_flags.py @@ -0,0 +1,48 @@ +from django.db import migrations + + +def copy_flags(apps, schema_editor): + CustomFlag = apps.get_model("redbox_core", "CustomFlag") + + try: + WaffleFlag = apps.get_model("waffle", "Flag") + except LookupError: + return + + for old in WaffleFlag.objects.all(): + new_flag, _ = CustomFlag.objects.get_or_create( + name=old.name, + defaults={ + "everyone": old.everyone, + "percent": old.percent, + "testing": old.testing, + "superusers": old.superusers, + "staff": old.staff, + "authenticated": old.authenticated, + "languages": old.languages, + "rollout": old.rollout, + "note": old.note, + "created": old.created, + "modified": old.modified, + }, + ) + + if old.groups.exists(): + new_flag.groups.set(old.groups.all()) + if old.users.exists(): + new_flag.users.set(old.users.all()) + + +def reverse_copy_flags(apps, schema_editor): + pass + + +class Migration(migrations.Migration): + + dependencies = [ + ("redbox_core", "0112_customflag"), + ] + + operations = [ + migrations.RunPython(copy_flags, reverse_copy_flags), + ] diff --git a/django_app/redbox_app/redbox_core/migrations/0114_remove_waffle_flag_table.py b/django_app/redbox_app/redbox_core/migrations/0114_remove_waffle_flag_table.py new file mode 100644 index 00000000..0721a1b9 --- /dev/null +++ b/django_app/redbox_app/redbox_core/migrations/0114_remove_waffle_flag_table.py @@ -0,0 +1,19 @@ +from django.db import migrations + + +class Migration(migrations.Migration): + + dependencies = [ + ("redbox_core", "0113_migrate_waffle_flags"), + ] + + operations = [ + migrations.RunSQL( + sql=""" + DROP TABLE IF EXISTS waffle_flag_users CASCADE; + DROP TABLE IF EXISTS waffle_flag_groups CASCADE; + DROP TABLE IF EXISTS waffle_flag CASCADE; + """, + reverse_sql=migrations.RunSQL.noop, + ), + ] diff --git a/django_app/redbox_app/redbox_core/utils.py b/django_app/redbox_app/redbox_core/utils.py index b7115be0..5736b73e 100644 --- a/django_app/redbox_app/redbox_core/utils.py +++ b/django_app/redbox_app/redbox_core/utils.py @@ -147,9 +147,9 @@ def user_has_invest_lens_access(request) -> bool: if waffle.flag_is_active(request, flag_name): return True - flag = waffle.get_waffle_flag_model().objects.get(name=flag_name) + flag = waffle.get_waffle_flag_model().objects.filter(name=flag_name).first() - if hasattr(flag, "get_extra_emails"): + if flag and hasattr(flag, "get_extra_emails"): user_email = getattr(request.user, "email", None) if user_email: extra_emails = flag.get_extra_emails() diff --git a/django_app/redbox_app/settings.py b/django_app/redbox_app/settings.py index 47592fc3..2c58003f 100644 --- a/django_app/redbox_app/settings.py +++ b/django_app/redbox_app/settings.py @@ -507,3 +507,4 @@ def filter_transactions(event, _hint): PRODUCT_NAME = env.str("PRODUCT_NAME", "Redbox at DBT") WAFFLE_FLAG_MODEL = "redbox_core.CustomFlag" +WAFFLE_CREATE_MISSING_FLAGS = True diff --git a/django_app/tests/conftest.py b/django_app/tests/conftest.py index 7942163e..b784951a 100644 --- a/django_app/tests/conftest.py +++ b/django_app/tests/conftest.py @@ -16,7 +16,9 @@ from django.utils import timezone from freezegun import freeze_time from moto import mock_aws +from waffle.models import get_waffle_flag_model +from redbox_app.redbox_core import flags from redbox_app.redbox_core.models import ( Agent, AISettings, @@ -118,14 +120,16 @@ def alice(create_user): @pytest.fixture def agents_list() -> list[Agent]: - agents = [ - Agent.objects.create( - name="Internal_Retrieval_Agent", - description="Fake", - agents_max_tokens=100, - llm_backend=ChatLLMBackend.objects.filter().first(), - ) - ] + agents = [] + + agent, _ = Agent.objects.get_or_create( + name="Internal_Retrieval_Agent", + description="Fake", + agents_max_tokens=100, + llm_backend=ChatLLMBackend.objects.first(), + ) + agents.append(agent) + for agent_name in [ "External_Retrieval_Agent", "Summarisation_Agent", @@ -137,7 +141,13 @@ def agents_list() -> list[Agent]: "Datahub_Agent", "Fake_Agent", ]: - agents += [Agent.objects.create(name=agent_name, description="Fake", agents_max_tokens=100)] + agent, _ = Agent.objects.get_or_create( + name=agent_name, + description="Fake", + agents_max_tokens=100, + ) + agents.append(agent) + return agents @@ -438,3 +448,37 @@ def default_agent() -> Agent: ) agent.save() return agent + + +CustomFlag = get_waffle_flag_model() + + +def ensure_critical_flags(): + flag_names = [ + "Internal_Retrieval_Agent", + ] + + for name in flag_names: + CustomFlag.objects.get_or_create( + name=name, + defaults={ + "everyone": False, + }, + ) + + +@pytest.fixture(autouse=True, scope="session") +def configure_waffle_for_tests(): + settings.WAFFLE_CREATE_MISSING_FLAGS = True + settings.WAFFLE_FLAG_MODEL = "redbox_core.CustomFlag" + + +@pytest.fixture(autouse=True) +def ensure_flags(): + CustomFlag.objects.get_or_create( + name=flags.ENABLE_INVEST_LENS, + defaults={ + "everyone": False, + "note": "test default", + }, + )