Improve tests

Author: pentest-gg

fastapi_users_db_dynamodb/__init__.py

@@ -231,7 +231,9 @@ async def delete(self, user: UP) -> None:
         try:
             await user.delete(condition=self.user_table.id.exists())  # type: ignore
         except DeleteError as e:
-            raise ValueError("User account could not be deleted.") from e
+            raise ValueError(  # pragma: no cover
+                "User account could not be deleted."
+            ) from e
         except PutError as e:
             if e.cause_response_code == "ConditionalCheckFailedException":
                 raise ValueError(

fastapi_users_db_dynamodb/access_token.py

@@ -127,7 +127,9 @@ async def delete(self, access_token: AP) -> None:
         try:
             await access_token.delete(condition=self.access_token_table.token.exists())  # type: ignore
         except DeleteError as e:
-            raise ValueError("Access token could not be deleted.") from e
+            raise ValueError(  # pragma: no cover
+                "Access token could not be deleted."
+            ) from e
         except PutError as e:
             if e.cause_response_code == "ConditionalCheckFailedException":
                 raise ValueError(

fastapi_users_db_dynamodb/tables.py

@@ -18,6 +18,9 @@ class CreatableTable(Protocol):
     @classmethod
     async def exists(cls) -> CoroutineType[Any, Any, bool] | bool: ...
 
+    @classmethod
+    async def delete_table(cls) -> CoroutineType[Any, Any, Any] | Any: ...
+
     @classmethod
     async def create_table(
         cls,
@@ -31,8 +34,14 @@ def __check_creatable_table(cls: type[Any]):
     """Check if an object is of type `CreatableTable`"""
     if not issubclass(cls, Model):
         raise TypeError(f"{cls.__name__} must be a subclass of Model")
-    if not hasattr(cls, "exists") or not hasattr(cls, "create_table"):
-        raise TypeError(f"{cls.__name__} must implement exists() and create_table()")
+    if (
+        not hasattr(cls, "exists")
+        or not hasattr(cls, "delete_table")
+        or not hasattr(cls, "create_table")
+    ):
+        raise TypeError(
+            f"{cls.__name__} must implement exists(), delete_table() and create_table()"
+        )
 
 
 async def ensure_tables_exist(*tables: type[CreatableTable]) -> None:
@@ -51,3 +60,17 @@ async def ensure_tables_exist(*tables: type[CreatableTable]) -> None:
                     wait=True,
                 )
             __tables_cache.add(table_cls)
+
+
+async def delete_tables(*tables: type[CreatableTable]) -> None:
+    """
+    Delete all given DynamoDB tables, if existent.
+    """
+    global __tables_cache
+
+    for table_cls in tables:
+        __check_creatable_table(table_cls)
+        if await table_cls.exists():
+            await table_cls.delete_table()
+        if table_cls in __tables_cache:
+            __tables_cache.remove(table_cls)

tests/test_access_token.py

@@ -98,11 +98,29 @@ async def test_queries(
     token_obj = await dynamodb_access_token_db.get_by_token("NOT_EXISTING_TOKEN")
     assert token_obj is None
 
+    # Create existing
+    with pytest.raises(
+        ValueError,
+        match="Access token could not be created because it already exists.",
+    ):
+        token = AccessToken()
+        token.token = "TOKEN"
+        token.user_id = user_id
+        await dynamodb_access_token_db.create(token)
+
     # Delete
     await dynamodb_access_token_db.delete(access_token)
     deleted_token = await dynamodb_access_token_db.get_by_token(access_token.token)
     assert deleted_token is None
 
+    # Update non-existent
+    new_time = now_utc()
+    with pytest.raises(
+        ValueError,
+        match="Access token could not be updated because it does not exist.",
+    ):
+        await dynamodb_access_token_db.update(access_token, {"created_at": new_time})
+
 
 @pytest.mark.asyncio
 async def test_insert_existing_token(

tests/test_misc.py

@@ -3,7 +3,7 @@
 
 from fastapi_users_db_dynamodb import config
 from fastapi_users_db_dynamodb.attributes import GUID
-from fastapi_users_db_dynamodb.tables import ensure_tables_exist
+from fastapi_users_db_dynamodb.tables import delete_tables, ensure_tables_exist
 
 
 class NotAModel:
@@ -14,6 +14,13 @@ class IncompleteModel(Model):
     pass
 
 
+class ValidModel(Model):
+    class Meta:
+        table_name: str = "valid_model_test"
+        region: str = config.get("DATABASE_REGION")
+        billing_mode: str = config.get("DATABASE_BILLING_MODE").value
+
+
 @pytest.mark.asyncio
 async def test_tables_invalid_models(monkeypatch):
     with pytest.raises(TypeError, match="must be a subclass of Model"):
@@ -22,6 +29,14 @@ async def test_tables_invalid_models(monkeypatch):
     with pytest.raises(AttributeError, match="PynamoDB Models require a"):
         await ensure_tables_exist(IncompleteModel)
 
+    with pytest.raises(AttributeError, match="PynamoDB Models require a"):
+        await delete_tables(IncompleteModel)
+
+    await ensure_tables_exist(ValidModel)
+    assert await ValidModel.exists()
+    await delete_tables(ValidModel)
+    assert not await ValidModel.exists()
+
     monkeypatch.delattr(Model, "exists", raising=True)
     with pytest.raises(TypeError):
         await ensure_tables_exist(IncompleteModel)