Added: ballot bdd to link votes together for 1 voter

Author: TheoSabattie

app/crud.py

@@ -218,9 +218,23 @@ def create_invite_tokens(
     _check_election_is_not_ended(get_election(db, election_ref))
     now = datetime.now()
     params = {"date_created": now, "date_modified": now, "election_ref": election_ref}
-    db_votes = [models.Vote(**params) for _ in range(num_voters * num_candidates)]
-    db.bulk_save_objects(db_votes, return_defaults=True)
-    db.commit()
+
+    try:
+        db_ballots = [models.Ballot(election_ref=election_ref) for _ in range(num_voters)]
+        db.bulk_save_objects(db_ballots, return_defaults=True)
+
+        db_votes = []
+
+        for ballot in db_ballots:
+            for _ in range(num_candidates):
+                db_votes.append(models.Vote(**params, ballot_id=ballot.id))
+        
+        db.bulk_save_objects(db_votes, return_defaults=True)
+        db.commit()
+    except Exception as e:
+        db.rollback()
+        raise e
+    
     vote_ids = [int(str(v.id)) for v in db_votes]
     tokens = [
         create_ballot_token(vote_ids[i::num_voters], election_ref)
@@ -405,14 +419,25 @@ def create_ballot(db: Session, ballot: schemas.BallotCreate) -> schemas.BallotGe
     )
     _check_ballot_is_consistent(election, ballot)
 
-    # Ideally, we would use RETURNING but it does not work yet for SQLite
-    db_votes = [
-        models.Vote(**v.model_dump(), election_ref=ballot.election_ref) for v in ballot.votes
-    ]
-    db.add_all(db_votes)
-    db.commit()
-    for v in db_votes:
-        db.refresh(v)
+    try:
+        db_ballot = models.Ballot(election_ref=ballot.election_ref)
+        db.add(db_ballot)
+        db.flush()
+
+        # Create votes and associate them with the ballot
+        db_votes = [
+            models.Vote(**v.model_dump(), election_ref=ballot.election_ref, ballot_id=db_ballot.id) 
+            for v in ballot.votes
+        ]
+        db.add_all(db_votes)
+        db.commit()
+        db.refresh(db_ballot)
+
+        for v in db_votes:
+            db.refresh(v)
+    except Exception as e:
+        db.rollback()
+        raise e
 
     votes_get = [schemas.VoteGet.model_validate(v) for v in db_votes]
     vote_ids = [v.id for v in votes_get]
@@ -511,6 +536,11 @@ def update_ballot(
     if len(db_votes) != len(vote_ids):
         raise errors.NotFoundError("votes")
 
+    # Verify all votes belong to the same ballot
+    ballot_ids = {v.ballot_id for v in db_votes if v.ballot_id is not None}
+    if len(ballot_ids) > 1:
+        raise errors.ForbiddenError("All votes must belong to the same ballot")
+
     election = schemas.ElectionGet.model_validate(db_votes[0].election)
 
     for vote, db_vote in zip(ballot.votes, db_votes):

app/models.py

@@ -2,7 +2,8 @@
 from sqlalchemy.sql import func
 from sqlalchemy.orm import relationship
 from .database import Base
-
+import uuid  # Pour pouvoir appeler uuid.uuid4()
+from sqlalchemy import UUID
 
 class Election(Base):
     __tablename__ = "elections"
@@ -24,6 +25,7 @@ class Election(Base):
     grades = relationship("Grade", back_populates="election")
     candidates = relationship("Candidate", back_populates="election")
     votes = relationship("Vote", back_populates="election")
+    ballots = relationship("Ballot", back_populates="election")
 
 
 class Candidate(Base):
@@ -72,3 +74,20 @@ class Vote(Base):
 
     election_ref = Column(String(20), ForeignKey("elections.ref"))
     election = relationship("Election", back_populates="votes")
+
+    ballot_id = Column(Integer, ForeignKey("ballots.id"), nullable=True) 
+    ballot = relationship("Ballot", back_populates="votes")
+
+
+class Ballot(Base):
+    __tablename__ = "ballots"
+
+    id = Column(Integer, primary_key=True, index=True)
+
+    voter_uuid = Column(UUID(as_uuid=True), default=uuid.uuid4, unique=True, index=True)
+
+    date_created = Column(DateTime, server_default=func.now())
+    election_ref = Column(String(20), ForeignKey("elections.ref"))
+
+    election = relationship("Election", back_populates="ballots")
+    votes = relationship("Vote", back_populates="ballot")