Added: ballot bdd to link votes together for 1 voter

Author: TheoSabattie

app/auth.py

@@ -27,12 +27,13 @@ def jws_verify(token: str) -> Mapping[str, t.Any]:
 def create_ballot_token(
     vote_ids: int | list[int],
     election_ref: str,
+    ballot_id:int
 ) -> str:
     if isinstance(vote_ids, int):
         vote_ids = [vote_ids]
     vote_ids = sorted(vote_ids)
     return jws.sign(
-        {"votes": vote_ids, "election": election_ref},
+        { "votes": vote_ids, "election": election_ref, "ballot": ballot_id },
         settings.secret,
         algorithm="HS256",
     )

app/crud.py

@@ -218,14 +218,33 @@ 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
+    
+    tokens = []
     vote_ids = [int(str(v.id)) for v in db_votes]
-    tokens = [
-        create_ballot_token(vote_ids[i::num_voters], election_ref)
-        for i in range(num_voters)
-    ]
+
+    for i, ballot in enumerate(db_ballots):
+        start = i * num_candidates
+        end = start + num_candidates
+        tokens.append(
+            create_ballot_token(vote_ids[start:end], election_ref, int(str(ballot.id)))
+        )
+
     return tokens
 
 
@@ -405,18 +424,29 @@ 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]
-    token = create_ballot_token(vote_ids, ballot.election_ref)
+    token = create_ballot_token(vote_ids, ballot.election_ref, db_ballot.id)
     return schemas.BallotGet(votes=votes_get, token=token, election=election)
 
 
@@ -511,6 +541,13 @@ 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:
+        print("ballots_ids:", sorted(ballot_ids))
+        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):
@@ -521,7 +558,7 @@ def update_ballot(
     db.commit()
 
     votes_get = [schemas.VoteGet.model_validate(v) for v in db_votes]
-    token = create_ballot_token(vote_ids, election_ref)
+    token = create_ballot_token(vote_ids, election_ref, db_votes[0].ballot_id)
     return schemas.BallotGet(votes=votes_get, token=token, election=election)
 
 

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")

app/tests/test_auth.py

@@ -41,9 +41,9 @@ def test_ballot_token():
     """
     vote_ids = list(range(1000))
     election_ref = "qwertyuiop"
-    token = create_ballot_token(vote_ids, election_ref)
+    token = create_ballot_token(vote_ids, election_ref, 1)
     data = jws_verify(token)
-    assert data == {"votes": vote_ids, "election": election_ref}
+    assert data == {"votes": vote_ids, "election": election_ref, "ballot": 1}
 
 
 def test_admin_token():