From bb924b4432cce6026221c7db093409b9ada2c0b2 Mon Sep 17 00:00:00 2001
From: Shreyas Sriram <31931776+shreyas-sriram@users.noreply.github.com>
Date: Tue, 11 May 2021 22:14:59 +0530
Subject: [PATCH] Add support for fee preference (#26)

* Add fee preference

* Advance fee check, use common function for preferences

* Update README.md

* Remove dev logs
---
 README.md |  1 +
 main.go   | 10 +++++++++-
 search.go | 20 +++++++++-----------
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index 5c8daa7..ee23d95 100644
--- a/README.md
+++ b/README.md
@@ -36,6 +36,7 @@ Flags:
   -a, --age int           Search appointment for age
   -d, --district string   Search by district name
   -e, --email string      Email address to send notifications
+  -f, --fee string        Fee preferences - free (or) paid. Default: No preference
   -h, --help              help for covaccine-notifier
   -i, --interval int      Interval to repeat the search. Default: (60) second
   -p, --password string   Email ID password for auth
diff --git a/main.go b/main.go
index 1d47c71..72ba963 100644
--- a/main.go
+++ b/main.go
@@ -12,7 +12,7 @@ import (
 )
 
 var (
-	pinCode, state, district, email, password, date, vaccine string
+	pinCode, state, district, email, password, date, vaccine, fee string
 
 	age, interval int
 
@@ -34,11 +34,15 @@ const (
 	emailPasswordEnv  = "EMAIL_PASSOWORD"
 	searchIntervalEnv = "SEARCH_INTERVAL"
 	vaccineEnv        = "VACCINE"
+	feeEnv            = "FEE"
 
 	defaultSearchInterval = 60
 
 	covishield = "covishield"
 	covaxin    = "covaxin"
+
+	free = "free"
+	paid = "paid"
 )
 
 func init() {
@@ -50,6 +54,7 @@ func init() {
 	rootCmd.PersistentFlags().StringVarP(&password, "password", "p", os.Getenv(emailPasswordEnv), "Email ID password for auth")
 	rootCmd.PersistentFlags().IntVarP(&interval, "interval", "i", getIntEnv(searchIntervalEnv), fmt.Sprintf("Interval to repeat the search. Default: (%v) second", defaultSearchInterval))
 	rootCmd.PersistentFlags().StringVarP(&vaccine, "vaccine", "v", os.Getenv(vaccineEnv), fmt.Sprintf("Vaccine preferences - covishield (or) covaxin. Default: No preference"))
+	rootCmd.PersistentFlags().StringVarP(&fee, "fee", "f", os.Getenv(feeEnv), fmt.Sprintf("Fee preferences - free (or) paid. Default: No preference"))
 }
 
 // Execute executes the main command
@@ -78,6 +83,9 @@ func checkFlags() error {
 	if !(vaccine == "" || vaccine == covishield || vaccine == covaxin) {
 		return errors.New("Invalid vaccine, please use covaxin or covishield")
 	}
+	if !(fee == "" || fee == free || fee == paid) {
+		return errors.New("Invalid fee preference, please use free or paid")
+	}
 	return nil
 }
 
diff --git a/search.go b/search.go
index 2d0a406..b10bf48 100644
--- a/search.go
+++ b/search.go
@@ -180,18 +180,13 @@ func searchByStateDistrict(age int, state, district string) error {
 	return getAvailableSessions(response, age)
 }
 
-// isPreferredVaccineAvailable checks for availability of preferred vaccine
-func isPreferredVaccineAvailable(current, preference string) bool {
-	switch preference {
-	case "":
+// isPreferredAvailable checks for availability of preferences
+func isPreferredAvailable(current, preference string) bool {
+	if preference == "" {
 		return true
-	case covishield:
-		return strings.ToLower(current) == covishield
-	case covaxin:
-		return strings.ToLower(current) == covaxin
+	} else {
+		return strings.ToLower(current) == preference
 	}
-
-	return false
 }
 
 func getAvailableSessions(response []byte, age int) error {
@@ -207,8 +202,11 @@ func getAvailableSessions(response []byte, age int) error {
 	var buf bytes.Buffer
 	w := tabwriter.NewWriter(&buf, 1, 8, 1, '\t', 0)
 	for _, center := range appnts.Centers {
+		if !isPreferredAvailable(center.FeeType, fee) {
+			continue
+		}
 		for _, s := range center.Sessions {
-			if s.MinAgeLimit <= age && s.AvailableCapacity != 0 && isPreferredVaccineAvailable(s.Vaccine, vaccine) {
+			if s.MinAgeLimit <= age && s.AvailableCapacity != 0 && isPreferredAvailable(s.Vaccine, vaccine) {
 				fmt.Fprintln(w, fmt.Sprintf("Center\t%s", center.Name))
 				fmt.Fprintln(w, fmt.Sprintf("State\t%s", center.StateName))
 				fmt.Fprintln(w, fmt.Sprintf("District\t%s", center.DistrictName))