Fix tests

Author: pentest-gg

fastapi_users_db_dynamodb/__init__.py

@@ -25,6 +25,7 @@
 from boto3.dynamodb.conditions import Attr
 from fastapi_users.db.base import BaseUserDatabase
 from fastapi_users.models import ID, OAP, UP
+from pydantic import BaseModel, ConfigDict, Field
 
 from fastapi_users_db_dynamodb._aioboto3_patch import *  # noqa: F403
 from fastapi_users_db_dynamodb.generics import GUID
@@ -34,22 +35,30 @@
 UUID_ID = uuid.UUID
 
 
-class DynamoDBBaseUserTable(Generic[ID]):
+class DynamoDBBaseUserTable(BaseModel, Generic[ID]):
     """Base user table schema for DynamoDB."""
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     __tablename__ = "user"
 
     if TYPE_CHECKING:
         id: ID
-    email: str
-    hashed_password: str
-    is_active: bool
-    is_superuser: bool
-    is_verified: bool
+    email: str = Field(..., description="The email of the user")
+    hashed_password: str = Field(..., description="The hashed password of the user")
+    is_active: bool = Field(
+        default=True, description="Whether the user is marked as active in the database"
+    )
+    is_superuser: bool = Field(
+        default=False, description="Whether the user has admin rights"
+    )
+    is_verified: bool = Field(
+        default=False, description="Whether the user has verified their email"
+    )
 
 
 class DynamoDBBaseUserTableUUID(DynamoDBBaseUserTable[UUID_ID]):
-    id: UUID_ID
+    id: UUID_ID = Field(default_factory=uuid.uuid4, description="The ID for the user")
 
 
 class DynamoDBBaseOAuthAccountTable(Generic[ID]):
@@ -59,18 +68,27 @@ class DynamoDBBaseOAuthAccountTable(Generic[ID]):
 
     if TYPE_CHECKING:
         id: ID
-    oauth_name: str
-    access_token: str
-    expires_at: int | None
-    refresh_token: str | None
-    account_id: str
-    account_email: str
-    user_id: ID
+    oauth_name: str = Field(..., description="The name of the OAuth social provider")
+    access_token: str = Field(
+        ..., description="The access token linked with the OAuth account"
+    )
+    expires_at: int | None = Field(
+        default=None, description="The timestamp at which this account expires"
+    )
+    refresh_token: str | None = Field(
+        default=None, description="The refresh token associated with this OAuth account"
+    )
+    account_id: str = Field(..., description="The ID of this OAuth account")
+    account_email: str = Field(
+        ..., description="The email associated with this OAuth account"
+    )
 
 
 class DynamoDBBaseOAuthAccountTableUUID(DynamoDBBaseUserTable[UUID_ID]):
-    id: UUID_ID
-    user_id: GUID
+    id: UUID_ID = Field(
+        default_factory=uuid.uuid4, description="The ID for the OAuth account"
+    )
+    user_id: GUID = Field(..., description="The user ID this OAuth account belongs to")
 
 
 class DynamoDBUserDatabase(Generic[UP, ID], BaseUserDatabase[UP, ID]):

fastapi_users_db_dynamodb/access_token.py

@@ -14,23 +14,30 @@
 import aioboto3
 from fastapi_users.authentication.strategy.db import AP, AccessTokenDatabase
 from fastapi_users.models import ID
+from pydantic import BaseModel, ConfigDict, Field
 
 from fastapi_users_db_dynamodb._aioboto3_patch import *  # noqa: F403
+from fastapi_users_db_dynamodb.generics import GUID
 
 
-class DynamoDBBaseAccessTokenTable(Generic[ID]):
+class DynamoDBBaseAccessTokenTable(BaseModel, Generic[ID]):
     """Base access token table schema for DynamoDB."""
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     __tablename__ = "accesstoken"
 
-    token: str
-    created_at: datetime
+    token: str = Field(..., description="The token value of the AccessToken object")
+    created_at: datetime = Field(
+        default_factory=lambda: datetime.now(timezone.utc),
+        description="The date of creation of the AccessToken object",
+    )
     if TYPE_CHECKING:
         user_id: ID
 
 
 class DynamoDBBaseAccessTokenTableUUID(DynamoDBBaseAccessTokenTable[uuid.UUID]):
-    user_id: uuid.UUID
+    user_id: GUID = Field(..., description="The user ID this token belongs to")
 
 
 class DynamoDBAccessTokenDatabase(Generic[AP], AccessTokenDatabase[AP]):

fastapi_users_db_dynamodb/generics.py

@@ -12,7 +12,7 @@
 from pydantic import UUID4
 
 
-class GUID:
+class GUID(UUID4):
     """
     Platform-independent GUID type.
 
@@ -49,7 +49,7 @@ def now_utc() -> datetime:
     return datetime.now(timezone.utc)
 
 
-class TIMESTAMPAware:
+class TIMESTAMPAware(datetime):
     """
     Kept for API compatibility.
 

tests/test_access_token.py

@@ -1,12 +1,11 @@
-import uuid
 from collections.abc import AsyncGenerator
 from datetime import datetime, timedelta, timezone
 
 import aioboto3
 import pytest
 import pytest_asyncio
 from moto import mock_aws
-from pydantic import UUID4
+from pydantic import UUID4, BaseModel
 
 from fastapi_users_db_dynamodb import DynamoDBBaseUserTableUUID, DynamoDBUserDatabase
 from fastapi_users_db_dynamodb._aioboto3_patch import *  # noqa: F403
@@ -18,7 +17,7 @@
 from tests.tables import ensure_table_exists
 
 
-class Base:
+class Base(BaseModel):
     pass
 
 
@@ -30,11 +29,6 @@ class User(DynamoDBBaseUserTableUUID, Base):
     pass
 
 
-@pytest.fixture
-def user_id() -> UUID4:
-    return uuid.uuid4()
-
-
 @pytest_asyncio.fixture
 async def dynamodb_access_token_db(
     user_id: UUID4,

tests/test_users.py

@@ -4,6 +4,7 @@
 import pytest
 import pytest_asyncio
 from moto import mock_aws
+from pydantic import BaseModel, ConfigDict, Field
 
 from fastapi_users_db_dynamodb import (
     UUID_ID,
@@ -16,12 +17,12 @@
 from tests.tables import ensure_table_exists
 
 
-class Base:
-    pass
+class Base(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
 
 class User(DynamoDBBaseUserTableUUID, Base):
-    first_name: str | None
+    first_name: str | None = Field(default=None, description="First name of the user")
 
 
 class OAuthBase:
@@ -33,8 +34,10 @@ class OAuthAccount(DynamoDBBaseOAuthAccountTableUUID, OAuthBase):
 
 
 class UserOAuth(DynamoDBBaseUserTableUUID, OAuthBase):
-    first_name: str | None
-    oauth_accounts: list[OAuthAccount]
+    first_name: str | None = Field(default=None, description="First name of the user")
+    oauth_accounts: list[OAuthAccount] = Field(
+        default_factory=list, description="Linked OAuth accounts"
+    )
 
 
 @pytest_asyncio.fixture
@@ -118,7 +121,7 @@ async def test_queries(dynamodb_user_db: DynamoDBUserDatabase[User, UUID_ID]):
         await dynamodb_user_db.get_by_oauth_account("foo", "bar")
     with pytest.raises(NotImplementedError):
         await dynamodb_user_db.add_oauth_account(user, {})
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(ValueError):
         oauth_account = OAuthAccount()
         await dynamodb_user_db.update_oauth_account(user, oauth_account, {})
 
@@ -161,23 +164,10 @@ async def test_queries_custom_fields(
 @pytest.mark.asyncio
 async def test_queries_oauth(
     dynamodb_user_db_oauth: DynamoDBUserDatabase[UserOAuth, UUID_ID],
+    oauth_account1,
+    oauth_account2,
 ):
     # Test OAuth accounts
-    oauth_account1 = {
-        "oauth_name": "service1",
-        "access_token": "TOKEN",
-        "expires_at": 1579000751,
-        "account_id": "user_oauth1",
-        "account_email": "king.arthur@camelot.bt",
-    }
-    oauth_account2 = {
-        "oauth_name": "service2",
-        "access_token": "TOKEN",
-        "expires_at": 1579000751,
-        "account_id": "user_oauth2",
-        "account_email": "king.arthur@camelot.bt",
-    }
-
     user_create = {"email": "lancelot@camelot.bt", "hashed_password": "guinevere"}
 
     # Create user