Now using http.Error() to send correct error code in header

handler/country.go

@@ -3,7 +3,6 @@ package handler
 import (
 	"assignment-1/pkg"
 	"encoding/json"
-	"fmt"
 	"github.com/gorilla/mux"
 	"log"
 	"net/http"
@@ -27,8 +26,7 @@ func Chandler(w http.ResponseWriter, r *http.Request) {
 			log.Fatal(err)
 		}
 	} else {
-		// need to check error for content before setting header
-		w.WriteHeader(http.StatusInternalServerError)
-		_, _ = fmt.Fprintln(w, "An error occured")
+		// need to check error for content when setting header (pkg.CorrectHeader)
+		http.Error(w, "An error occured", pkg.CorrectHeader(err))
 	}
 }

handler/diag.go

@@ -3,7 +3,6 @@ package handler
 import (
 	"assignment-1/pkg"
 	"encoding/json"
-	"fmt"
 	"log"
 	"net/http"
 )
@@ -12,8 +11,7 @@ func Dhandler(w http.ResponseWriter, r *http.Request) {
 	var diagnostics pkg.Diag
 
 	if err := pkg.GetDiag(&diagnostics); err != nil {
-		w.WriteHeader(http.StatusBadGateway)
-		_, _ = fmt.Fprintln(w, "An error occured")
+		http.Error(w, "An error occured", http.StatusBadGateway)
 	}
 
 	w.Header().Set("Content-Type", "application/json")
@@ -21,4 +19,3 @@ func Dhandler(w http.ResponseWriter, r *http.Request) {
 		log.Fatal(err)
 	}
 }
-

handler/species.go

@@ -20,9 +20,7 @@ func Shandler(w http.ResponseWriter, r *http.Request) {
 			log.Fatal(err)
 		}
 	} else {
-		w.WriteHeader(http.StatusBadGateway)
+		w.WriteHeader(pkg.CorrectHeader(err))
 		_, _ = fmt.Fprintln(w, "An error occured")
 	}
 }
-
-

pkg/country.go

@@ -12,6 +12,7 @@ type Country struct {
 }
 
 type tempCountry struct {
+	Code string `json:"alpha2code"`
 	Flag string `json:"flag"`
 	Name string `json:"name"`
 }
@@ -32,7 +33,6 @@ func GetCountryByCode(code, limit string) (Country, error) {
 	var tmpC tempCountry
 	var species response
 	urlCountry := countryApi + code
-	urlSpecies := occurrenceApi + "country=" + code + "&limit=" + limit
 
 	// gets country info to temp struct
 	err := getBody(urlCountry, &tmpC)
@@ -43,8 +43,9 @@ func GetCountryByCode(code, limit string) (Country, error) {
 	// add tmp info to country struct
 	country.CountryFlag = tmpC.Flag
 	country.CountryName = tmpC.Name
-	country.Code        = code			// was passed with the func call
+	country.Code = tmpC.Code // was passed with the func call
 
+	urlSpecies := occurrenceApi + "country=" + country.Code + "&limit=" + limit
 	// gets species in specified country
 	err = getBody(urlSpecies, &species)
 	if err != nil {
@@ -67,6 +68,3 @@ func GetCountryByCode(code, limit string) (Country, error) {
 
 	return country, nil
 }
-
-
-

pkg/funcs.go

@@ -6,7 +6,6 @@ import (
 	"net/http"
 )
 
-
 // allowing access to "overloaded" functions
 func unfold(m mashup, response *http.Response) error {
 	return m.unmarshal(response)
@@ -35,7 +34,6 @@ func (s *Specie) unmarshal(resp *http.Response) error {
 	return err
 }
 
-
 // parses httpresponse as a string used later for parsing
 func getBody(url string, m mashup) error {
 	resp, err := http.Get(url)
@@ -47,7 +45,21 @@ func getBody(url string, m mashup) error {
 		return err
 	}
 	if resp.StatusCode != http.StatusOK {
-		return fmt.Errorf("status not 200 OK, got: %d", resp.StatusCode)
+		return fmt.Errorf("%d", resp.StatusCode)
 	}
 	return nil
 }
+
+// func for returning correct http.Head response code
+func CorrectHeader(err error) int {
+
+	switch err.Error() {
+	case "400", "404":
+		return http.StatusBadGateway
+	case "408":
+		return http.StatusGatewayTimeout
+	default:
+		return http.StatusInternalServerError
+	}
+
+}