fix

app/routers/finish.go

@@ -10,20 +10,18 @@ type FinishRequestBody struct {
 	Id string `json:"id"`
 }
 
-func Finish(w http.ResponseWriter, r *http.Request) {
+func Finish(r *http.Request) (interface{}, int) {
 	d := json.NewDecoder(r.Body)
 	body := FinishRequestBody{}
 	err := d.Decode(&body)
 	if err != nil {
-		w.WriteHeader(http.StatusBadRequest)
-		return
+		return nil, http.StatusBadRequest
 	}
 
 	err = tasks.Finish(body.Id)
 	if err != nil {
-		w.WriteHeader(http.StatusInternalServerError)
-		return
+		return nil, http.StatusInternalServerError
 	}
 
-	w.WriteHeader(http.StatusAccepted)
+	return nil, http.StatusAccepted
 }

app/routers/put.go

@@ -15,13 +15,12 @@ type PutRequestBody struct {
 	Delay            *int   `json:"delay"`
 }
 
-func Put(w http.ResponseWriter, r *http.Request) {
+func Put(r *http.Request) (interface{}, int) {
 	d := json.NewDecoder(r.Body)
 	body := PutRequestBody{}
 	err := d.Decode(&body)
 	if err != nil {
-		w.WriteHeader(http.StatusBadRequest)
-		return
+		return nil, http.StatusBadRequest
 	}
 
 	queue := r.Header.Get("queue")
@@ -44,9 +43,8 @@ func Put(w http.ResponseWriter, r *http.Request) {
 
 	err = tasks.Add(task)
 	if err != nil {
-		w.WriteHeader(http.StatusInternalServerError)
-		return
+		return nil, http.StatusInternalServerError
 	}
 
-	w.WriteHeader(http.StatusAccepted)
+	return nil, http.StatusAccepted
 }

app/routers/take.go

@@ -1,7 +1,6 @@
 package routers
 
 import (
-	"encoding/json"
 	"net/http"
 	tasks "queues-go/app/storage/mongo/collections"
 )
@@ -16,13 +15,11 @@ type TakeResponse struct {
 	Task *TaskResponse `json:"task"`
 }
 
-func Take(w http.ResponseWriter, r *http.Request) {
+func Take(r *http.Request) (interface{}, int) {
 	queue := r.Header.Get("queue")
 	task, err := tasks.Take(queue)
 	if err != nil {
-		println("Error taking")
-		w.WriteHeader(http.StatusInternalServerError)
-		return
+		return nil, http.StatusInternalServerError
 	}
 
 	var response TakeResponse
@@ -36,12 +33,5 @@ func Take(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
-	data, err := json.Marshal(response)
-	if err != nil {
-		println("Error marshalling")
-		w.WriteHeader(http.StatusInternalServerError)
-		return
-	}
-
-	w.Write(data)
+	return response, http.StatusOK
 }

main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"encoding/json"
 	"log"
 	"net/http"
 	"queues-go/app/routers"
@@ -8,11 +9,22 @@ import (
 	"time"
 )
 
-func handlerWrapper(f func(http.ResponseWriter, *http.Request)) func(http.ResponseWriter, *http.Request) {
+func handlerWrapper(f func(*http.Request) (interface{}, int)) func(http.ResponseWriter, *http.Request) {
 	return func(w http.ResponseWriter, r *http.Request) {
 		start := time.Now()
-		f(w, r)
-		log.Printf("%s %s", r.URL, time.Since(start))
+		response, status := f(r)
+		if response != nil {
+			data, err := json.Marshal(response)
+			if err != nil {
+				log.Printf("Error marshalling response")
+				status = http.StatusInternalServerError
+			}
+			w.WriteHeader(status)
+			w.Write(data)
+		} else {
+			w.WriteHeader(status)
+		}
+		log.Printf("%s %d %s", r.URL, status, time.Since(start))
 	}
 }