Merge pull request #801 from i-dot-ai/bugfix/REDBOX-533-return-404-wh…

…ere-appropriate

Return 404 where appropriate.

django_app/redbox_app/redbox_core/views.py

@@ -189,7 +189,7 @@ def ingest_file(uploaded_file: UploadedFile, user: User) -> Sequence[str]:
 
 @login_required
 def remove_doc_view(request, doc_id: uuid):
-    file = File.objects.get(pk=doc_id)
+    file = get_object_or_404(File, id=doc_id)
     errors: list[str] = []
 
     if request.method == "POST":
@@ -224,7 +224,7 @@ def get(self, request: HttpRequest, chat_id: uuid.UUID | None = None) -> HttpRes
         messages: Sequence[ChatMessage] = []
         current_chat = None
         if chat_id:
-            current_chat = ChatHistory.objects.get(id=chat_id)
+            current_chat = get_object_or_404(ChatHistory, id=chat_id)
             if current_chat.users != request.user:
                 return redirect(reverse("chats"))
             messages = (
@@ -294,7 +294,7 @@ def post(self, request: HttpRequest, chat_id: uuid.UUID) -> HttpResponse:
 class CitationsView(View):
     @method_decorator(login_required)
     def get(self, request: HttpRequest, message_id: uuid.UUID | None = None) -> HttpResponse:
-        message = ChatMessage.objects.get(id=message_id)
+        message = get_object_or_404(ChatMessage, id=message_id)
 
         if message.chat_history.users != request.user:
             return redirect(reverse("chats"))
@@ -403,7 +403,7 @@ def file_status_api_view(request: HttpRequest) -> JsonResponse:
         logger.error("Error getting file object information - no file ID provided %s.")
         return JsonResponse({"status": StatusEnum.unknown.label})
     try:
-        file = File.objects.get(pk=file_id)
+        file = get_object_or_404(File, id=file_id)
     except File.DoesNotExist as ex:
         logger.exception("File object information not found in django - file does not exist %s.", file_id, exc_info=ex)
         return JsonResponse({"status": StatusEnum.unknown.label})

django_app/tests/test_views.py

@@ -227,6 +227,20 @@ def test_remove_doc_view(client: Client, alice: User, file_pdf_path: Path, s3_cl
         assert File.objects.get(id=new_file.id).status == StatusEnum.deleted
 
 
+@pytest.mark.django_db()
+def test_remove_nonexistent_doc(alice: User, client: Client):
+    # Given
+    client.force_login(alice)
+    nonexistent_uuid = uuid.uuid4()
+
+    # When
+    url = reverse("remove-doc", kwargs={"doc_id": nonexistent_uuid})
+    response = client.get(url)
+
+    # Then
+    assert response.status_code == HTTPStatus.NOT_FOUND
+
+
 @pytest.mark.django_db()
 def test_post_message_to_new_session(alice: User, client: Client, requests_mock: Mocker):
     # Given
@@ -373,6 +387,20 @@ def test_view_session_with_documents(chat_message: ChatMessage, client: Client):
     assert b"original_file.txt" in response.content
 
 
+@pytest.mark.django_db()
+def test_nonexistent_chats(alice: User, client: Client):
+    # Given
+    client.force_login(alice)
+    nonexistent_uuid = uuid.uuid4()
+
+    # When
+    url = reverse("chats", kwargs={"chat_id": nonexistent_uuid})
+    response = client.get(url)
+
+    # Then
+    assert response.status_code == HTTPStatus.NOT_FOUND
+
+
 @pytest.mark.django_db()
 def test_post_chat_title(alice: User, chat_history: ChatHistory, client: Client):
     # Given
@@ -508,6 +536,20 @@ def test_user_cannot_see_other_users_citations(chat_message_with_citation: ChatH
     assert response.headers.get("Location") == "/chats/"
 
 
+@pytest.mark.django_db()
+def test_nonexistent_citations(alice: User, client: Client):
+    # Given
+    client.force_login(alice)
+    nonexistent_uuid = uuid.uuid4()
+
+    # When
+    url = reverse("citations", kwargs={"message_id": nonexistent_uuid})
+    response = client.get(url)
+
+    # Then
+    assert response.status_code == HTTPStatus.NOT_FOUND
+
+
 @pytest.mark.django_db()
 def test_staff_user_can_see_route(chat_history_with_files: ChatHistory, client: Client):
     # Given