✨ Add back the user instance to the authentication system, block the addition of oneself as a friend, implement a DTO for the friends list and modify the architecture of the Friend model to distinguish the two users.

app/dto/FriendAddDTO.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+
+class FriendAddDTO(BaseModel):
+    friend_user_id: str

app/dto/FriendListDTO.py

@@ -0,0 +1,8 @@
+from pydantic import BaseModel
+
+from app.models.friend import Friend
+
+
+class FriendListDTO(BaseModel):
+    friends: list[Friend]
+    external_friends: list[Friend]

app/dto/__init__.py

@@ -1 +1,3 @@
-from .UserRegisterDTO import UserRegisterDTO
+from .UserRegisterDTO import UserRegisterDTO
+from .FriendAddDTO import FriendAddDTO
+from .FriendListDTO import FriendListDTO

app/main.py

@@ -49,7 +49,7 @@ def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
 
     return encoded_jwt
 
-async def get_current_user(token: str = Depends(oauth2_scheme)):
+async def get_current_user(token: str = Depends(oauth2_scheme)) -> User:
     credentials_exception = HTTPException(
         status_code=status.HTTP_401_UNAUTHORIZED,
         detail="Could not validate credentials",
@@ -69,7 +69,7 @@ async def get_current_user(token: str = Depends(oauth2_scheme)):
     if user is None:
         raise credentials_exception
     
-    return user
+    return user_serialize(user)
 
 
 # Routes - TODO: find workaround to display 401/409/... HTTP error codes in openapi.json 
@@ -88,7 +88,7 @@ async def register(user: UserRegisterDTO):
     access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
     access_token = create_access_token(data={"sub": user.username}, expires_delta=access_token_expires)
 
-    return {"access_token": access_token, "token_type": "bearer"}
+    return {"access_token": access_token, "token_type": "bearer", "user_id": str(user["_id"])}
 
 @app.post("/login", response_model=Token)
 async def login(form_data: OAuth2PasswordRequestForm = Depends()):
@@ -103,7 +103,7 @@ async def login(form_data: OAuth2PasswordRequestForm = Depends()):
     access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
     access_token = create_access_token(data={"sub": form_data.username}, expires_delta=access_token_expires)
 
-    return {"access_token": access_token, "token_type": "bearer"}
+    return {"access_token": access_token, "token_type": "bearer", "user_id": str(user["_id"])}
 
 @app.get("/logout")
 async def logout(current_user: User = Depends(get_current_user)):
@@ -145,9 +145,20 @@ async def get_friend(id: str, current_user: User = Depends(get_current_user)):
     return friend
 
 @app.post("/friend/add")
-async def add_friend(friend: Friend, current_user: User = Depends(get_current_user)):
+async def add_friend(friendAdd: FriendAddDTO, current_user: User = Depends(get_current_user)):
     # TODO: test if exists
-    friend_id = friends_collection.insert_one(friend.model_dump()).inserted_id
+    friend: Friend = friendAdd.model_dump()
+
+    print(current_user)
+
+    if(current_user.uid == friend["friend_user_id"]):
+        raise HTTPException(
+            status_code=status.HTTP_409_CONFLICT,
+            detail="Cannot add yourself as a friend"
+        )
+
+    friend["user_id"] = current_user.uid
+    friend_id = friends_collection.insert_one(friend).inserted_id
     return {"id": str(friend_id)}
 
 @app.delete("/friend/{id}/delete")
@@ -174,10 +185,9 @@ async def deny_friend(id: str, current_user: User = Depends(get_current_user)):
     
     return {"message": "Friend request denied"}
 
-@app.get("/friends")
+@app.get("/friends", response_model=FriendListDTO)
 async def list_friends(current_user: User = Depends(get_current_user)):
-    friends = friends_serialize(friends_collection.find().to_list())
-    return friends
+    return FriendListDTO(**friends_serialize(friends_collection.find({"user_id": current_user.uid}).to_list(), friends_collection.find({"friend_user_id": current_user.uid})))
 
 @app.get("/users")
 async def search_users(name: str, current_user: User = Depends(get_current_user)):

app/models/friend.py

@@ -2,4 +2,6 @@ from pydantic import BaseModel
 
 
 class Friend(BaseModel):
-    user_id: str
+    user_id: str
+    friend_user_id: str
+    status: str

app/models/token.py

@@ -2,4 +2,5 @@ from pydantic import BaseModel
 
 class Token(BaseModel):
     access_token: str
-    token_type: str
+    token_type: str
+    user_id: str