From a1c6b284151ead42faf4e9c1e8fd15e246c25cd6 Mon Sep 17 00:00:00 2001 From: George Burton Date: Thu, 12 Sep 2024 06:50:00 +0100 Subject: [PATCH] use get_user_model everywhere --- django_app/redbox_app/redbox_core/consumers.py | 3 ++- django_app/redbox_app/redbox_core/serializers.py | 5 ++++- django_app/redbox_app/redbox_core/views/document_views.py | 4 +++- django_app/tests/conftest.py | 4 +++- django_app/tests/management/test_commands.py | 6 +++++- django_app/tests/test_admin.py | 5 ++++- django_app/tests/test_consumers.py | 5 ++++- django_app/tests/views/test_chat_views.py | 4 +++- django_app/tests/views/test_citation_views.py | 4 +++- django_app/tests/views/test_document_views.py | 5 ++++- django_app/tests/views/test_ratings_views.py | 5 ++++- 11 files changed, 39 insertions(+), 11 deletions(-) diff --git a/django_app/redbox_app/redbox_core/consumers.py b/django_app/redbox_app/redbox_core/consumers.py index 493856ba2..35a4330ac 100644 --- a/django_app/redbox_app/redbox_core/consumers.py +++ b/django_app/redbox_app/redbox_core/consumers.py @@ -8,6 +8,7 @@ from channels.db import database_sync_to_async from channels.generic.websocket import AsyncWebsocketConsumer from django.conf import settings +from django.contrib.auth import get_user_model from django.forms.models import model_to_dict from django.utils import timezone from langchain_core.documents import Document @@ -26,9 +27,9 @@ ChatRoleEnum, Citation, File, - User, ) +User = get_user_model() OptFileSeq = Sequence[File] | None logger = logging.getLogger(__name__) logger.info("WEBSOCKET_SCHEME is: %s", settings.WEBSOCKET_SCHEME) diff --git a/django_app/redbox_app/redbox_core/serializers.py b/django_app/redbox_app/redbox_core/serializers.py index f34d14a44..7304f0cfe 100644 --- a/django_app/redbox_app/redbox_core/serializers.py +++ b/django_app/redbox_app/redbox_core/serializers.py @@ -1,6 +1,9 @@ +from django.contrib.auth import get_user_model from rest_framework import serializers -from redbox_app.redbox_core.models import Chat, ChatMessage, ChatMessageTokenUse, File, User +from redbox_app.redbox_core.models import Chat, ChatMessage, ChatMessageTokenUse, File + +User = get_user_model() class FileSerializer(serializers.ModelSerializer): diff --git a/django_app/redbox_app/redbox_core/views/document_views.py b/django_app/redbox_app/redbox_core/views/document_views.py index 8390c5e59..bb04a21b4 100644 --- a/django_app/redbox_app/redbox_core/views/document_views.py +++ b/django_app/redbox_app/redbox_core/views/document_views.py @@ -3,6 +3,7 @@ from collections.abc import MutableSequence, Sequence from pathlib import Path +from django.contrib.auth import get_user_model from django.contrib.auth.decorators import login_required from django.core.exceptions import FieldError, ValidationError from django.core.files.uploadedfile import UploadedFile @@ -15,9 +16,10 @@ from django_q.tasks import async_task from requests.exceptions import RequestException -from redbox_app.redbox_core.models import File, StatusEnum, User +from redbox_app.redbox_core.models import File, StatusEnum from redbox_app.worker import ingest +User = get_user_model() logger = logging.getLogger(__name__) CHUNK_SIZE = 1024 # move this somewhere diff --git a/django_app/tests/conftest.py b/django_app/tests/conftest.py index cd7898ed0..07357a82e 100644 --- a/django_app/tests/conftest.py +++ b/django_app/tests/conftest.py @@ -8,6 +8,7 @@ import pytest from botocore.exceptions import ClientError from django.conf import settings +from django.contrib.auth import get_user_model from django.core.files.uploadedfile import SimpleUploadedFile, UploadedFile from django.core.management import call_command from django.utils import timezone @@ -22,9 +23,10 @@ Citation, File, StatusEnum, - User, ) +User = get_user_model() + logger = logging.getLogger(__name__) diff --git a/django_app/tests/management/test_commands.py b/django_app/tests/management/test_commands.py index 9bc2c1d0e..914fbba19 100644 --- a/django_app/tests/management/test_commands.py +++ b/django_app/tests/management/test_commands.py @@ -6,13 +6,17 @@ import pytest from botocore.exceptions import UnknownClientMethodError from django.conf import settings +from django.contrib.auth import get_user_model from django.core.management import CommandError, call_command from django.utils import timezone from freezegun import freeze_time from magic_link.models import MagicLink from requests_mock import Mocker -from redbox_app.redbox_core.models import Chat, ChatMessage, ChatRoleEnum, File, StatusEnum, User +from redbox_app.redbox_core.models import Chat, ChatMessage, ChatRoleEnum, File, StatusEnum + +User = get_user_model() + # === check_file_status command tests === diff --git a/django_app/tests/test_admin.py b/django_app/tests/test_admin.py index 303b77495..206a201c4 100644 --- a/django_app/tests/test_admin.py +++ b/django_app/tests/test_admin.py @@ -6,13 +6,16 @@ import pytest from bs4 import BeautifulSoup +from django.contrib.auth import get_user_model from django.test import Client from django.urls import reverse from yarl import URL -from redbox_app.redbox_core.models import ChatMessage, User +from redbox_app.redbox_core.models import ChatMessage from redbox_app.redbox_core.serializers import ChatMessageSerializer, ChatSerializer, UserSerializer +User = get_user_model() + logger = logging.getLogger(__name__) diff --git a/django_app/tests/test_consumers.py b/django_app/tests/test_consumers.py index 7fbb4ec07..0874a1ef6 100644 --- a/django_app/tests/test_consumers.py +++ b/django_app/tests/test_consumers.py @@ -9,6 +9,7 @@ import pytest from channels.db import database_sync_to_async from channels.testing import WebsocketCommunicator +from django.contrib.auth import get_user_model from django.db.models import Model from langchain_core.documents import Document from langchain_core.language_models import BaseChatModel @@ -20,9 +21,11 @@ from redbox.models.chat import MetadataDetail from redbox_app.redbox_core import error_messages from redbox_app.redbox_core.consumers import ChatConsumer -from redbox_app.redbox_core.models import Chat, ChatMessage, ChatMessageTokenUse, ChatRoleEnum, File, User +from redbox_app.redbox_core.models import Chat, ChatMessage, ChatMessageTokenUse, ChatRoleEnum, File from redbox_app.redbox_core.prompts import CHAT_MAP_QUESTION_PROMPT +User = get_user_model() + logging.basicConfig(level=os.environ.get("LOG_LEVEL", "INFO")) logger = logging.getLogger(__name__) diff --git a/django_app/tests/views/test_chat_views.py b/django_app/tests/views/test_chat_views.py index d6e1f2d73..833f4f4d1 100644 --- a/django_app/tests/views/test_chat_views.py +++ b/django_app/tests/views/test_chat_views.py @@ -5,15 +5,17 @@ import pytest from bs4 import BeautifulSoup +from django.contrib.auth import get_user_model from django.test import Client from django.urls import reverse from redbox_app.redbox_core.models import ( Chat, ChatMessage, - User, ) +User = get_user_model() + logger = logging.getLogger(__name__) diff --git a/django_app/tests/views/test_citation_views.py b/django_app/tests/views/test_citation_views.py index e7fead5ae..49640e1b1 100644 --- a/django_app/tests/views/test_citation_views.py +++ b/django_app/tests/views/test_citation_views.py @@ -5,6 +5,7 @@ import pytest from bs4 import BeautifulSoup +from django.contrib.auth import get_user_model from django.test import Client from django.urls import reverse @@ -14,9 +15,10 @@ ChatRoleEnum, Citation, File, - User, ) +User = get_user_model() + logger = logging.getLogger(__name__) diff --git a/django_app/tests/views/test_document_views.py b/django_app/tests/views/test_document_views.py index 9230a155e..5ab010c8f 100644 --- a/django_app/tests/views/test_document_views.py +++ b/django_app/tests/views/test_document_views.py @@ -6,10 +6,13 @@ import pytest from botocore.exceptions import ClientError from django.conf import settings +from django.contrib.auth import get_user_model from django.test import Client from django.urls import reverse -from redbox_app.redbox_core.models import File, StatusEnum, User +from redbox_app.redbox_core.models import File, StatusEnum + +User = get_user_model() logger = logging.getLogger(__name__) diff --git a/django_app/tests/views/test_ratings_views.py b/django_app/tests/views/test_ratings_views.py index 93a11aae4..ae0ec0201 100644 --- a/django_app/tests/views/test_ratings_views.py +++ b/django_app/tests/views/test_ratings_views.py @@ -3,10 +3,13 @@ from http import HTTPStatus import pytest +from django.contrib.auth import get_user_model from django.test import Client from django.urls import reverse -from redbox_app.redbox_core.models import ChatMessage, User +from redbox_app.redbox_core.models import ChatMessage + +User = get_user_model() logger = logging.getLogger(__name__)