Adding more tests for adding a poll if it doesn't exist on get, Adding code to update votes

server/services/poll.go

@@ -3,6 +3,7 @@ package services
 import (
 	"database/sql"
 	"errors"
+	"time"
 
 	"go-sjles-pta-vote/server/db"
 	"go-sjles-pta-vote/server/models"
@@ -10,6 +11,7 @@ import (
 
 var ErrQuestionAlreadyExists = errors.New("Question already exists")
 var ErrQuestionDoesntExist = errors.New("Question does not exist yet")
+var ErrVoterAlreadyVoted = errors.New("Voter already voted")
 
 func CreatePoll(poll *models.Poll) (*models.Poll, error) {
 	new_poll := models.Poll{}
@@ -125,4 +127,110 @@ func GetPollByQuestion(question string) (*models.Poll, error) {
 	}
 
 	return &new_poll, nil
+}
+
+func GetAndCreatePollByQuestion(question string) (*models.Poll, error) {
+	new_poll, err := GetPollByQuestion(question)
+
+	if err == ErrQuestionDoesntExist {
+		create_poll := &models.Poll{
+			Question: question,
+			ExpiresAt: time.Now().Add(time.Hour * 10).Format("2006-01-02 15:04:05"),
+		}
+
+		if _, err = CreatePoll(create_poll); err != nil {
+			return nil, err
+		}
+
+		return GetPollByQuestion(question)
+	} else if err != nil {
+		return nil, err
+	} else {
+		return new_poll, err
+	}
+}
+
+func SetVote(poll_id int64, email string, vote bool) error {
+	db_conn, err := db.Connect()
+	if err != nil {
+		return err
+	}
+	defer db.Close()
+
+	set_voter_stmt, err := db_conn.Prepare(`
+		INSERT IGNORE INTO voters
+			(poll_id, voter_email)
+		VALUES ($1, $2)
+	`)
+	if err != nil {
+		return err
+	}
+	defer set_voter_stmt.Close()
+
+	res, err := set_voter_stmt.Exec(poll_id, email)
+	if err != nil {
+		return err
+	} else {
+		rows_changed, err := res.RowsAffected()
+		if rows_changed != 1 {
+			return ErrVoterAlreadyVoted
+		} else if err != nil {
+			return err
+		}
+	}
+
+	is_voter_member_stmt, err := db_conn.Prepare(`
+		SELECT 1
+		FROM members
+		WHERE email == $1
+	`)
+	if err != nil {
+		return err
+	}
+	defer is_voter_member_stmt.Close()
+
+	var member_check int64
+	is_member := true
+	err = is_voter_member_stmt.QueryRow(email).Scan(&member_check)
+	if err == sql.ErrNoRows {
+		is_member = false
+	} else if err != nil {
+		return err
+	}
+
+	// Member column name is not dependant on user input
+	//	So it's ok to put it directly in the query
+	member_column_name := "member_"
+	if !is_member {
+		member_column_name = "non_" + member_column_name
+	}
+
+	if vote {
+		member_column_name += "yes_votes"
+	} else {
+		member_column_name += "no_votes"
+	}
+
+	add_vote_stmt, err := db_conn.Prepare(`
+		UPDATE polls
+		SET ` + member_column_name + ` = ` + member_column_name + ` 1
+		WHERE id == $1
+	`)
+	if err != nil {
+		return err
+	}
+	defer add_vote_stmt.Close()
+
+	res, err = add_vote_stmt.Exec(poll_id)
+	if err != nil {
+		return err
+	}
+
+	if num, err := res.RowsAffected(); num != 1 {
+		return errors.New("Failed to update votes")
+	} else if err != nil {
+		return err
+	}
+
+	return nil
 }