diff --git a/main.go b/main.go index 2c72330..2f85522 100644 --- a/main.go +++ b/main.go @@ -86,7 +86,7 @@ func (t *Test) collectData() { return // really needed? } -func (t *Test) startTest(w http.ResponseWriter, req *http.Request) { +func (t *Test) Start(w http.ResponseWriter, req *http.Request) { // Check if another collectData is running (better: check sync.Mutex) if t.Status == "ONGOING" { http.Error(w, "Another test is currently ongoing, stop or try again later", 409) //better error code @@ -144,7 +144,7 @@ func (t *Test) startTest(w http.ResponseWriter, req *http.Request) { go t.collectData() } -func (t *Test) stopTest(w http.ResponseWriter, req *http.Request) { +func (t *Test) Stop(w http.ResponseWriter, req *http.Request) { if t.Status != "ONGOING" { e := fmt.Sprintf("Cannot stop test in %s status: try start first", t.Status) http.Error(w, e, 400) @@ -157,18 +157,54 @@ func (t *Test) stopTest(w http.ResponseWriter, req *http.Request) { fmt.Fprintf(w, "%v\n", t.Result) } -func (t *Test) testStatus(w http.ResponseWriter, req *http.Request) { +func (t *Test) getStatus(w http.ResponseWriter, req *http.Request) { + qStream := req.URL.Query().Get("stream") qFormat := req.URL.Query().Get("format") - if qFormat == "json" { - b, err := json.Marshal(t.Result) - if err != nil { - log.Printf("Cannot marshal json: %s\n", err) - http.Error(w, err.Error(), 500) + + flusher, ok := w.(http.Flusher) + if !ok { + if qStream == "true" { + log.Printf("Http server does not support flusher\n") + http.Error(w, "Server does not support flusher", + http.StatusInternalServerError) return } - fmt.Fprintf(w, "%s\n", b) + } + + if qStream == "true" { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + } + + if qFormat == "json" { + for { + b, err := json.Marshal(t.Result) + if err != nil { + log.Printf("Cannot marshal json: %s\n", err) + http.Error(w, err.Error(), 500) + return + } + + fmt.Fprintf(w, "%s\n", b) + + if qStream == "true" { + time.Sleep(time.Duration(t.Rate) * time.Millisecond) + flusher.Flush() + } else { + return + } + } } else { - fmt.Fprintf(w, "%v\n", t.Result) + for { + fmt.Fprintf(w, "%+ v\n", t.Result) + if qStream == "true" { + time.Sleep(time.Duration(t.Rate) * time.Millisecond) + flusher.Flush() + } else { + return + } + } fmt.Fprintf(w, "Test %s, status: %s\n", t.Name, t.Status) } } @@ -193,9 +229,9 @@ func main() { t := newTest() log.Printf("%+v\n", t) - http.HandleFunc("/start", t.startTest) - http.HandleFunc("/status", t.testStatus) - http.HandleFunc("/stop", t.stopTest) + http.HandleFunc("/start", t.Start) + http.HandleFunc("/status", t.getStatus) + http.HandleFunc("/stop", t.Stop) log.Printf("Listening on %s\n", *listen) log.Fatal(http.ListenAndServe(*listen, nil))