feat: add budgets

Makefile

@@ -11,7 +11,7 @@ lint:
 
 test:
 	rm test.db 2> /dev/null || true
-	DB_NAME="./test.db" poetry run pytest -sv .
+	DB_NAME="./test.db" poetry run pytest -svv .
 
 run: install
 	poetry run uvicorn src.main:app --host localhost --port 8000 --reload

src/crud.py

@@ -73,3 +73,41 @@ def get_spendings_by_group_id(db: Session, group_id: int):
         .limit(100)
         .all()
     )
+
+
+################################################
+# BUDGETS
+################################################
+
+
+def create_budget(db: Session, budget: schemas.BudgetCreate):
+    db_budget = models.Budget(**dict(budget))
+    db.add(db_budget)
+    db.commit()
+    db.refresh(db_budget)
+    return db_budget
+
+
+def get_budget_by_id(db: Session, budget_id: int):
+    return db.query(models.Budget).filter(models.Budget.id == budget_id).first()
+
+
+def put_budget(db: Session, budget_id: int, budget: schemas.BudgetPut):
+    db_budget = db.query(models.Budget).filter(models.Budget.id == budget_id).first()
+    db_budget.amount = budget.amount
+    db_budget.description = budget.description
+    db_budget.start_date = budget.start_date
+    db_budget.end_date = budget.end_date
+    db_budget.category_id = budget.category_id
+    db.commit()
+    db.refresh(db_budget)
+    return db_budget
+
+
+def get_budgets_by_group_id(db: Session, group_id: int):
+    return (
+        db.query(models.Budget)
+        .filter(models.Budget.group_id == group_id)
+        .limit(100)
+        .all()
+    )

src/main.py

@@ -129,3 +129,38 @@ def list_group_spendings(db: DbDependency, user: UserDependency, group_id: int):
             status_code=HTTPStatus.NOT_FOUND, detail="Grupo inexistente"
         )
     return crud.get_spendings_by_group_id(db, group_id)
+
+
+################################################
+# BUDGETS
+################################################
+
+
+@app.post("/budget", status_code=HTTPStatus.CREATED)
+def create_budget(
+    spending: schemas.BudgetCreate, db: DbDependency, user: UserDependency
+):
+    # TODO: check group exists
+    return crud.create_budget(db, spending)
+
+
+@app.get("/budget/{budget_id}")
+def get_budget(db: DbDependency, user: UserDependency, budget_id: int):
+    return crud.get_budget_by_id(db, budget_id)
+
+
+@app.put("/budget/{budget_id}")
+def put_budget(
+    db: DbDependency, user: UserDependency, budget_id: int, budget: schemas.BudgetPut
+):
+    return crud.put_budget(db, budget_id, budget)
+
+
+@app.get("/group/{group_id}/budget")
+def list_group_budgets(db: DbDependency, user: UserDependency, group_id: int):
+    group = crud.get_group_by_id(db, group_id)
+    if group is None or group.owner_id != user.id:
+        return HTTPException(
+            status_code=HTTPStatus.NOT_FOUND, detail="Grupo inexistente"
+        )
+    return crud.get_budgets_by_group_id(db, group_id)

src/models.py

@@ -31,3 +31,15 @@ class Spending(Base):
     amount = Column(Integer)
     description = Column(String)
     date: Mapped[datetime] = mapped_column(DateTime, default=func.now())
+
+
+class Budget(Base):
+    __tablename__ = "budgets"
+
+    id = Column(Integer, primary_key=True)
+    group_id = Column(ForeignKey("groups.id"))
+    category_id = Column(Integer)  # TODO: Column(ForeignKey("categories.id"))
+    start_date: Mapped[datetime] = mapped_column(DateTime)
+    end_date: Mapped[datetime] = mapped_column(DateTime)
+    amount = Column(Integer)
+    description = Column(String)

src/schemas.py

@@ -1,4 +1,4 @@
-import datetime
+from datetime import datetime
 from typing import Optional
 from pydantic import BaseModel, Field
 
@@ -51,7 +51,7 @@ class Group(GroupBase):
 class SpendingBase(BaseModel):
     amount: int
     description: str
-    date: Optional[datetime.date] = Field(None)
+    date: Optional[datetime] = Field(None)
     group_id: int
 
 
@@ -62,3 +62,29 @@ class SpendingCreate(SpendingBase):
 class Spending(SpendingBase):
     id: int
     owner_id: int
+
+
+################################################
+# BUDGETS
+################################################
+
+
+class BudgetBase(BaseModel):
+    amount: int
+    description: str
+    start_date: datetime
+    end_date: datetime
+    category_id: int
+
+
+class BudgetCreate(BudgetBase):
+    group_id: int
+
+
+class BudgetPut(BudgetBase):
+    pass
+
+
+class Budget(BudgetBase):
+    id: int
+    group_id: int

src/test_main.py

@@ -257,3 +257,83 @@ def test_get_spendings(
     assert response.status_code == HTTPStatus.OK
     assert len(response.json()) == 1
     assert schemas.Spending(**response.json()[0]) == some_spending
+
+
+################################################
+# BUDGETS
+################################################
+
+
+@pytest.fixture
+def some_budget(client: TestClient, some_user_id: int, some_group: schemas.Group):
+    response = client.post(
+        url="/budget",
+        json={
+            "amount": 500,
+            "description": "café",
+            "start_date": "2021-01-01",
+            "end_date": "2021-02-01",
+            "group_id": some_group.id,
+            "category_id": 1,
+        },
+        headers={"x-user": str(some_user_id)},
+    )
+
+    assert response.status_code == HTTPStatus.CREATED
+    response_body = response.json()
+    assert "id" in response_body
+    assert response_body["group_id"] == some_group.id
+    return schemas.Budget(**response_body)
+
+
+def test_create_new_budget(client: TestClient, some_budget: schemas.Budget):
+    # NOTE: test is inside fixture
+    pass
+
+
+def test_get_budget(client: TestClient, some_user_id: int, some_budget: schemas.Budget):
+    response = client.get(
+        url=f"/budget/{some_budget.id}",
+        headers={"x-user": str(some_user_id)},
+    )
+
+    assert response.status_code == HTTPStatus.OK
+    assert schemas.Budget(**response.json()) == some_budget
+
+
+def test_put_budget(client: TestClient, some_user_id: int, some_budget: schemas.Budget):
+    put_body = {
+        "amount": 1000,
+        "description": "some other description",
+        "start_date": "2021-03-01T00:00:00",
+        "end_date": "2021-04-01T00:00:00",
+        "category_id": 2,
+    }
+    response = client.put(
+        url=f"/budget/{some_budget.id}",
+        json=put_body,
+        headers={"x-user": str(some_user_id)},
+    )
+
+    assert response.status_code == HTTPStatus.OK
+
+    response = client.get(
+        url=f"/budget/{some_budget.id}",
+        headers={"x-user": str(some_user_id)},
+    )
+    assert response.status_code == HTTPStatus.OK
+    for k, v in put_body.items():
+        assert response.json()[k] == v
+
+
+def test_get_group_budgets(
+    client: TestClient, some_user_id: int, some_budget: schemas.Budget
+):
+    response = client.get(
+        url=f"/group/{some_budget.group_id}/budget",
+        headers={"x-user": str(some_user_id)},
+    )
+
+    assert response.status_code == HTTPStatus.OK
+    assert len(response.json()) == 1
+    assert schemas.Budget(**response.json()[0]) == some_budget