add test and authentication for message service

alembic/migrations/env.py

@@ -22,7 +22,7 @@ if config.config_file_name is not None:
 # for 'autogenerate' support
 # from myapp import mymodel
 # target_metadata = mymodel.Base.metadata
-target_metadata = None
+target_metadata = Base.metadata
 
 # other values from the config, defined by the needs of env.py,
 # can be acquired:

core/api/__init__.py

@@ -1,3 +1,11 @@
 from fastapi import APIRouter
+from .message import router as message_router
+from .chat import router as chat_router
+from .auth import router as auth_router
 
 router = APIRouter(prefix='/api/v1')
+
+# Include routers
+router.include_router(message_router)
+router.include_router(chat_router)
+router.include_router(auth_router)

core/api/auth/__init__.py

@@ -0,0 +1 @@
+from .handlers import router

core/api/auth/handlers.py

@@ -0,0 +1,12 @@
+from fastapi import APIRouter
+
+from core.models.auth.requests import LoginRequest
+from core.models.auth.responses import LoginResponse
+from core.services import auth_service
+router = APIRouter(prefix='/auth', tags=['auth'])
+
+
+@router.post("")
+async def login(user: LoginRequest) -> LoginResponse:
+
+    return await auth_service.login_user(user.username)

core/api/chat/__init__.py

@@ -0,0 +1 @@
+from .handlers import router

core/api/chat/handlers.py

@@ -0,0 +1,15 @@
+from fastapi import APIRouter, Depends, Response
+
+from core.helpers.auth.helpers import get_current_user
+from core.models.message.db import MPProfile
+from core.services import chat_service
+
+from core.models.message.requests import CreateChatRequest
+
+router = APIRouter(prefix='/chat', tags=['chat'])
+
+
+@router.post("")
+async def create_chat(response: Response, chat: CreateChatRequest, user: MPProfile = Depends(get_current_user)):
+    response.status_code = 201
+    return await chat_service.create_chat(chat.name, user.id)

core/api/message/__init__.py

@@ -0,0 +1 @@
+from .handlers import router

core/api/message/handlers.py

@@ -0,0 +1,5 @@
+from fastapi import APIRouter, Depends, HTTPException, status
+
+router = APIRouter(prefix="/message", tags=["message"])
+
+

core/errors/errors.py

@@ -0,0 +1,27 @@
+from core.models.wrapper.responses import Response, Meta
+from fastapi.responses import JSONResponse
+
+
+class CustomExceptionHandler(Exception):
+    def __init__(self, status_code: int, response: Response):
+        self.status_code = status_code
+        self.response = response
+
+    def result(self) -> JSONResponse:
+        return JSONResponse(status_code=self.status_code,
+                            content=self.response.model_dump(exclude_none=True))
+
+
+def not_authenticated_error():
+    response = Response(status=401,
+                        error=True,
+                        message="Token is not valid")
+    raise CustomExceptionHandler(status_code=401, response=response)
+
+
+def incorrect_login():
+    response = Response(status=401,
+                        error=True,
+                        message="Incorrect login")
+    raise CustomExceptionHandler(status_code=401, response=response)
+

core/helpers/auth/helpers.py

@@ -0,0 +1,21 @@
+import jwt
+from jwt.exceptions import PyJWTError
+from config import Config
+from core.models.message.db import MPProfile
+from fastapi import Header
+from core.storage import auth_storage
+from core.errors.errors import not_authenticated_error
+
+
+async def get_current_user(token: str = Header(alias="Authorization")) -> MPProfile:
+    """Get the current user."""
+    try:
+        token = token.split("Bearer ")[1]
+        token_data = jwt.decode(token, Config.secret, algorithms=["HS256"])
+        user = await auth_storage.get_user_by_id(token_data["user_id"])
+        if user:
+            return user
+        raise not_authenticated_error()
+
+    except PyJWTError:
+        not_authenticated_error()

core/models/auth/requests.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+
+class LoginRequest(BaseModel):
+    username: int

core/models/message/db.py

@@ -0,0 +1,38 @@
+from sqlalchemy import Column, Integer, ForeignKey
+from sqlalchemy.orm import mapped_column, Mapped, relationship
+
+from database import Base
+
+
+class MPProfile(Base):
+    __tablename__ = 'mp_profile'
+
+    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+    external_id: Mapped[int] = mapped_column(unique=True)
+    chats: Mapped[list['MPChat']] = relationship('MPChat', secondary='mp_chat_user')
+
+
+class MPChat(Base):
+    __tablename__ = 'mp_chat'
+    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+    name: Mapped[str | None]
+    admin_id: Mapped[int] = mapped_column(ForeignKey('mp_profile.id'))
+    admin: Mapped[MPProfile] = relationship()
+    users: Mapped[list[MPProfile]] = relationship('MPProfile', secondary='mp_chat_user')
+
+
+class MPMessage(Base):
+    __tablename__ = 'mp_message'
+    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+    sender_id: Mapped[int] = mapped_column(ForeignKey('mp_profile.id'))
+    chat_id: Mapped[int] = mapped_column(ForeignKey('mp_chat.id'))
+    content: Mapped[str]
+    chat: Mapped[MPChat] = relationship()
+    sender: Mapped[MPProfile] = relationship()
+
+
+class MPChatUser(Base):
+    __tablename__ = 'mp_chat_user'
+    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+    user_id: Mapped[int] = mapped_column(ForeignKey('mp_profile.id', ondelete='CASCADE'))
+    chat_id: Mapped[int] = mapped_column(ForeignKey('mp_chat.id', ondelete='CASCADE'))

core/models/message/requests.py

@@ -0,0 +1,11 @@
+from pydantic import BaseModel
+
+
+class CreateChatRequest(BaseModel):
+    name: str
+
+
+class SendMessageRequest(BaseModel):
+    chat_id: int
+    sender_id: int
+    content: str

core/models/message/responses.py

@@ -0,0 +1,3 @@
+from pydantic import BaseModel
+
+

core/models/wrapper/responses.py

@@ -0,0 +1,19 @@
+import datetime
+
+from pydantic import BaseModel
+
+
+class Meta(BaseModel):
+    request_started: datetime.datetime
+    request_finished: datetime.datetime
+
+
+class Response(BaseModel):
+    status: int
+    error: bool
+    message: str
+    data: BaseModel | None = None
+    meta: Meta | None = None
+
+
+

core/services/__init__.py

@@ -0,0 +1,7 @@
+from .auth.services import Service as AuthService
+from .chat.services import Service as ChatService
+
+# Register services
+
+auth_service = AuthService()
+chat_service = ChatService()

core/services/auth/services.py

@@ -0,0 +1,19 @@
+import datetime
+
+from config import Config
+from core.models.auth.responses import LoginResponse
+from core.models.message.db import MPProfile
+from core.storage import auth_storage
+from core.errors.errors import incorrect_login
+import jwt
+
+
+class Service:
+    async def login_user(self, username: str) -> LoginResponse:
+        user = await auth_storage.get_user(external_id=int(username))
+        if not user:
+            incorrect_login()
+        token = jwt.encode({'user_id': user.id,
+                    "exp": (datetime.datetime.now() + datetime.timedelta(minutes=Config.token_lifetime)).timestamp()},
+                   Config.secret, algorithm='HS256')
+        return LoginResponse(access_token=token)

core/services/chat/services.py

@@ -0,0 +1,9 @@
+from core.models.message.db import MPChat
+from core.storage import chat_storage
+
+
+class Service:
+    async def create_chat(self, name: str, admin_id: int) -> MPChat:
+        chat = await chat_storage.create_chat(name=name, admin_id=admin_id)
+        await chat_storage.add_member(chat.id, admin_id)
+        return chat

core/storage/__init__.py

@@ -2,4 +2,9 @@
 Register our storages
 """
 
+from database import Session
+from .chat.storage import Storage as ChatStorage
+from .auth.storage import Storage as AuthStorage
 
+chat_storage = ChatStorage(Session)
+auth_storage = AuthStorage(Session)

core/storage/auth/storage.py

@@ -0,0 +1,37 @@
+from sqlalchemy import select, Result
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from ..base import BaseStorage
+from ...models.message.db import MPProfile
+
+
+class Storage(BaseStorage):
+    async def create_user(self, external_id: int) -> int:
+        """
+        Storage handler for creating a user
+        """
+        async with self.get_session() as session:
+            user = MPProfile(external_id=external_id)
+            session.add(user)
+            await session.commit()
+            return user.id
+
+    async def get_user(self, external_id: int) -> MPProfile | None:
+        """
+        Storage handler for getting a user
+        """
+        async with self.get_session() as session:
+            session: AsyncSession
+            query = select(MPProfile).where(MPProfile.external_id == external_id)
+            result: Result = await session.execute(query)
+            user: MPProfile | None = result.scalar_one_or_none()
+            return user
+
+    async def get_user_by_id(self, id: int) -> MPProfile | None:
+        """
+        Storage handler for getting a user by id
+        """
+        async with self.get_session() as session:
+            session: AsyncSession
+            user = await session.get(MPProfile, id)
+            return user

core/storage/base.py

@@ -14,4 +14,4 @@ class BaseStorage:
         try:
             yield session
         finally:
-            await session.aclose()
+            await session.aclose()

core/storage/chat/storage.py

@@ -0,0 +1,39 @@
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from ..base import BaseStorage
+from core.models.message.db import MPChat, MPProfile
+
+
+class Storage(BaseStorage):
+    async def create_chat(self, name: str, admin_id: int) -> MPChat:
+        """
+        Storage handler for creating a chat
+        """
+        async with self.get_session() as session:
+            chat = MPChat(name=name, admin_id=admin_id)
+            session.add(chat)
+            await session.commit()
+            await session.refresh(chat)
+            return chat
+
+    async def add_member(self, chat_id: int, user_id: int):
+        """
+        Storage handler for adding a member to a chat
+        """
+        async with self.get_session() as session:
+            session: AsyncSession
+            chat = await session.get(MPChat, chat_id)
+            user: MPProfile = await session.get(MPProfile, user_id)
+            await session.refresh(chat, attribute_names=["users"])
+            chat.users.append(user)
+            await session.commit()
+            await session.refresh(chat, attribute_names=["users", "admin"])
+
+
+    async def get_chat(self, id: int):
+        """
+        Storage handler for getting a chat
+        """
+        async with self.get_session() as session:
+            chat = await session.get(MPChat, id)
+            return chat

database.py

@@ -1,17 +1,31 @@
+import datetime
 import json
 
 import asyncpg
-from sqlalchemy.orm import DeclarativeBase
+from sqlalchemy import func
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
 from sqlalchemy.ext.asyncio import AsyncAttrs, create_async_engine, async_sessionmaker
 from config import Config
 from loguru import logger
 
 
 class Base(AsyncAttrs, DeclarativeBase):
-    pass
+    created_at: Mapped[datetime.datetime] = mapped_column(server_default=func.now())
+    modified_at: Mapped[datetime.datetime] = mapped_column(server_default=func.now(), onupdate=func.now())
 
 
-engine = create_async_engine(str(Config.postgres_url), pool_size=20, max_overflow=0)
+class Engine:
+    _engine = None
 
-Session = async_sessionmaker(engine, expire_on_commit=False)
+    @property
+    def engine(self):
+        return self._engine
+
+    def __init__(self):
+        if self._engine is None:
+            self._engine = create_async_engine(str(Config.postgres_url), pool_size=20, max_overflow=0)
+
+
+engine = Engine()
+Session = async_sessionmaker(engine.engine, expire_on_commit=False)
 

docker-compose.yml

@@ -0,0 +1,11 @@
+services:
+  database:
+    image: postgres:alpine
+    environment:
+      POSTGRES_DB: mp_message
+      POSTGRES_USER: mp_user
+      POSTGRES_PASSWORD: mp_password
+    ports:
+      - "5432:5432"
+    volumes:
+      - ./data:/var/lib/postgresql/data

main.py

@@ -8,6 +8,7 @@ from starlette.middleware.cors import CORSMiddleware
 from config import Config
 
 from core.api import router as api_router
+from core.errors.errors import CustomExceptionHandler
 
 app = FastAPI()
 
@@ -33,5 +34,10 @@ async def on_startup():
     logger.success('Application startup complete at {time}', time=datetime.now(tz=timezone.utc))
 
 
+@app.exception_handler(CustomExceptionHandler)
+async def custom_exception_handler(_, exc):
+    return exc.result()
+
+
 if __name__ == '__main__':
     uvicorn.run('main:app', host=str(Config.host), port=Config.port, reload=True)

requirements.in

@@ -2,8 +2,11 @@ fastapi
 uvicorn
 sqlalchemy[asyncio]
 asyncpg
-python-camelcaser
 loguru
 pydantic-settings
 pydantic[email]
-alembic
+alembic
+pytest
+httpx
+pyjwt
+trio

requirements.txt

@@ -1,5 +1,5 @@
 #
-# This file is autogenerated by pip-compile with Python 3.10
+# This file is autogenerated by pip-compile with Python 3.12
 # by the following command:
 #
 #    pip-compile
@@ -11,35 +11,56 @@ annotated-types==0.6.0
 anyio==3.7.1
     # via
     #   fastapi
+    #   httpx
     #   starlette
-async-timeout==4.0.3
-    # via asyncpg
 asyncpg==0.29.0
     # via -r requirements.in
+attrs==23.2.0
+    # via
+    #   outcome
+    #   trio
+certifi==2024.2.2
+    # via
+    #   httpcore
+    #   httpx
 click==8.1.7
     # via uvicorn
 dnspython==2.4.2
     # via email-validator
 email-validator==2.1.0.post1
     # via pydantic
-exceptiongroup==1.2.0
-    # via anyio
 fastapi==0.104.1
     # via -r requirements.in
 greenlet==3.0.2
     # via sqlalchemy
 h11==0.14.0
-    # via uvicorn
+    # via
+    #   httpcore
+    #   uvicorn
+httpcore==1.0.5
+    # via httpx
+httpx==0.27.0
+    # via -r requirements.in
 idna==3.6
     # via
     #   anyio
     #   email-validator
+    #   httpx
+    #   trio
+iniconfig==2.0.0
+    # via pytest
 loguru==0.7.2
     # via -r requirements.in
 mako==1.3.0
     # via alembic
 markupsafe==2.1.3
     # via mako
+outcome==1.3.0.post0
+    # via trio
+packaging==24.0
+    # via pytest
+pluggy==1.5.0
+    # via pytest
 pydantic[email]==2.5.2
     # via
     #   -r requirements.in
@@ -49,20 +70,27 @@ pydantic-core==2.14.5
     # via pydantic
 pydantic-settings==2.1.0
     # via -r requirements.in
-pyenchant==3.2.2
+pyjwt==2.8.0
-    # via python-camelcaser
+    # via -r requirements.in
-python-camelcaser==1.0.2
+pytest==8.2.0
     # via -r requirements.in
 python-dotenv==1.0.0
     # via pydantic-settings
 sniffio==1.3.0
-    # via anyio
+    # via
+    #   anyio
+    #   httpx
+    #   trio
+sortedcontainers==2.4.0
+    # via trio
 sqlalchemy[asyncio]==2.0.23
     # via
     #   -r requirements.in
     #   alembic
 starlette==0.27.0
     # via fastapi
+trio==0.25.1
+    # via -r requirements.in
 typing-extensions==4.9.0
     # via
     #   alembic
@@ -70,6 +98,5 @@ typing-extensions==4.9.0
     #   pydantic
     #   pydantic-core
     #   sqlalchemy
-    #   uvicorn
 uvicorn==0.24.0.post1
     # via -r requirements.in

tests/test_auth.py

@@ -0,0 +1,34 @@
+import threading
+
+import jwt
+import uvicorn
+
+from config import Config
+from main import app
+from httpx import Client
+
+server_thread = threading.Thread(target=uvicorn.run, args=(app,), daemon=True,
+                                 kwargs={"host": "127.0.0.1", "port": 8000, "reload": False})
+server_thread.start()
+
+
+client = Client(base_url="http://127.0.0.1:8000")
+
+
+def test_auth_no_username():
+    resp = client.post('/api/v1/auth')
+    assert resp.status_code == 422
+
+
+def test_auth_incorrect_username():
+    resp = client.post('/api/v1/auth', json={'username': -1})
+    assert resp.status_code == 401
+
+
+def test_auth_correct_username():
+    resp = client.post('/api/v1/auth', json={'username': 1})
+    assert resp.status_code == 200
+    assert 'access_token' in resp.json()
+    assert len(resp.json()['access_token']) > 0
+    data = jwt.decode(resp.json()['access_token'], Config.secret, algorithms=['HS256'])
+    assert data["user_id"] == 2

tests/test_chat.py

@@ -0,0 +1,31 @@
+import threading
+
+import jwt
+import uvicorn
+
+from config import Config
+from main import app
+from httpx import Client
+
+server_thread = threading.Thread(target=uvicorn.run, args=(app,), daemon=True,
+                                 kwargs={"host": "127.0.0.1", "port": 8000, "reload": False})
+server_thread.start()
+
+
+client = Client(base_url="http://127.0.0.1:8000")
+
+
+def test_create_chat():
+    # Complete authorization
+    resp = client.post('/api/v1/auth', json={'username': 1})
+    assert resp.status_code == 200
+    token = resp.json()['access_token']
+
+    resp = client.post('/api/v1/chat', json={'name': 'test_chat'}, headers={'Authorization': f'Bearer {token}'})
+    print(resp.json())
+    assert resp.status_code == 201
+    assert 'id' in resp.json()
+    assert resp.json()['id'] > 0
+    assert 'name' in resp.json()
+    assert resp.json()['name'] == 'test_chat'
+