chore: Protect chatbot view using feature flag (#2151)

chore: Protect chatbot view using feature flag (#2151)

* Protect chatbot view using feature flag

* Protect stream_chat_response view

* Add migration that creates "chatbot" flag

* Adjust tests

---------

Co-authored-by: Pedro Costa <pedrobeethoven8@gmail.com>
diff --git a/ankihub/ai/migrations/0010_auto_20240612_1624.py b/ankihub/ai/migrations/0010_auto_20240612_1624.py
new file mode 100755
index 0000000..111463f
--- /dev/null
+++ b/ankihub/ai/migrations/0010_auto_20240612_1624.py
@@ -0,0 +1,26 @@
+# Generated by Django 4.2.8 on 2024-06-12 16:24
+
+from django.db import migrations
+
+flag_name = "chatbot"
+
+
+def create_flag(apps, schema_editor):
+    Flag = apps.get_model("waffle", "Flag")
+    Flag.objects.get_or_create(name=flag_name)
+
+
+def delete_flag(apps, schema_editor):
+    Flag = apps.get_model("waffle", "Flag")
+    Flag.objects.filter(name=flag_name).delete()
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("ai", "0009_aisearch_results_amount"),
+        ("waffle", "0001_initial"),
+    ]
+
+    operations = [
+        migrations.RunPython(create_flag, delete_flag),
+    ]
diff --git a/ankihub/ai/tests/test_views.py b/ankihub/ai/tests/test_views.py
index c384945..bbd45ed 100644
--- a/ankihub/ai/tests/test_views.py
+++ b/ankihub/ai/tests/test_views.py
@@ -9,6 +9,7 @@ from django.db.models import Case, FloatField, Value, When
 from django.test import Client
 from django.urls import reverse
 from langchain.schema import Document
+from waffle.testutils import override_flag
 
 from ankihub.ai.models import NoteEmbedding
 from ankihub.ai.tests.factories import NoteEmbeddingFactory
@@ -969,6 +970,7 @@ class TestUpdateModalDisplayUserSettingView:
         )
 
 
+@override_flag("chatbot", active=True)
 def test_chatbot(user):
     url = reverse("ai:chatbot")
     client = Client()
@@ -980,6 +982,17 @@ def test_chatbot(user):
     assert "ai/chatbot.html" in (t.name for t in response.templates)
 
 
+def test_chatbot_without_feature_flag_return_404(user):
+    url = reverse("ai:chatbot")
+    client = Client()
+    client.force_login(user=user)
+
+    response = client.get(url)
+
+    assert response.status_code == 404
+
+
+@override_flag("chatbot", active=True)
 @patch("ankihub.ai.views.llm_conversation_stream")
 def test_stream_chatbot_response(mock_llm_conversation_stream, user):
     class LLMResponse:
@@ -1011,6 +1024,15 @@ def test_stream_chatbot_response(mock_llm_conversation_stream, user):
     assert messages == [answer.content for answer in stream_response]
 
 
+def test_stream_chatbot_response_without_feature_flag_returns_404(user):
+    url = reverse("ai:stream_chatbot_response")
+    client = Client()
+    client.force_login(user=user)
+    response = client.post(url)
+    assert response.status_code == 404
+
+
+@override_flag("chatbot", active=True)
 @patch("ankihub.ai.views.llm_conversation_stream")
 def test_stream_chatbot_with_invalid_question(mock_llm_conversation_stream, user):
     url = reverse("ai:stream_chatbot_response")
diff --git a/ankihub/ai/views.py b/ankihub/ai/views.py
index f497372..7c36de9 100644
--- a/ankihub/ai/views.py
+++ b/ankihub/ai/views.py
@@ -7,6 +7,7 @@ from django.conf import settings
 from django.contrib.auth.decorators import login_required
 from django.core.paginator import EmptyPage, Paginator
 from django.http import (
+    Http404,
     HttpResponse,
     HttpResponseBadRequest,
     JsonResponse,
@@ -18,6 +19,7 @@ from django.urls import reverse
 from django.views.decorators.http import require_http_methods
 from langchain_community.document_loaders.pdf import AmazonTextractPDFLoader
 from tld.exceptions import TldBadUrl, TldDomainNotFound, TldIOError
+from waffle.decorators import flag_is_active
 from youtube_transcript_api import NoTranscriptFound
 
 from ankihub.ai.forms import FileUploadForm, LinkSearchForm
@@ -465,12 +467,17 @@ def update_modal_display_user_setting(request):
 
 @knox_token_or_login_required
 def chatbot(request):
+    if not flag_is_active(request, "chatbot"):
+        raise Http404()
     return render(request, "ai/chatbot.html")
 
 
 @knox_token_or_login_required
 @require_http_methods(["POST"])
 def stream_chatbot_response(request):
+    if not flag_is_active(request, "chatbot"):
+        raise Http404()
+
     user_question = request.POST.get("question")
 
     if not user_question:

GitHub
sha: e707c49c2565236d991ffd5c1c7406d61cd1695f