From 3dc7b02b6dd98c596119ba07a7928d630dc61173 Mon Sep 17 00:00:00 2001
From: Georges Haidar <ghaidar0@gmail.com>
Date: Mon, 15 Jan 2024 15:53:34 +0000
Subject: [PATCH] chore: add event streaming test endpoints (#5)

This change adds mock endpoints that stream server-sent events. These will support the test suite in openapi-generation.

Also added the ability to run the server using a different bind address. This is helpful when running the project locally outside of docker.
---
 cmd/server/main.go               |  20 ++++-
 go.mod                           |   2 +-
 internal/eventstreams/service.go | 128 +++++++++++++++++++++++++++++++
 3 files changed, 147 insertions(+), 3 deletions(-)
 create mode 100644 internal/eventstreams/service.go

diff --git a/cmd/server/main.go b/cmd/server/main.go
index 6682d37..d64501e 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -1,11 +1,13 @@
 package main
 
 import (
+	"flag"
 	"log"
 	"net/http"
 
 	"github.com/speakeasy-api/speakeasy-api-test-service/internal/acceptHeaders"
 	"github.com/speakeasy-api/speakeasy-api-test-service/internal/errors"
+	"github.com/speakeasy-api/speakeasy-api-test-service/internal/eventstreams"
 	"github.com/speakeasy-api/speakeasy-api-test-service/internal/pagination"
 	"github.com/speakeasy-api/speakeasy-api-test-service/internal/readonlywriteonly"
 	"github.com/speakeasy-api/speakeasy-api-test-service/internal/responseHeaders"
@@ -16,7 +18,11 @@ import (
 	"github.com/speakeasy-api/speakeasy-api-test-service/internal/requestbody"
 )
 
+var bindArg = flag.String("b", ":8080", "Bind address")
+
 func main() {
+	flag.Parse()
+
 	r := mux.NewRouter()
 	r.HandleFunc("/ping", func(w http.ResponseWriter, r *http.Request) {
 		_, _ = w.Write([]byte("pong"))
@@ -33,9 +39,19 @@ func main() {
 	r.HandleFunc("/readonlyorwriteonly", readonlywriteonly.HandleReadOrWrite).Methods(http.MethodPost)
 	r.HandleFunc("/readonlyandwriteonly", readonlywriteonly.HandleReadAndWrite).Methods(http.MethodPost)
 	r.HandleFunc("/writeonlyoutput", readonlywriteonly.HandleWriteOnlyOutput).Methods(http.MethodPost)
+	r.HandleFunc("/eventstreams/json", eventstreams.HandleEventStreamJSON).Methods(http.MethodPost)
+	r.HandleFunc("/eventstreams/text", eventstreams.HandleEventStreamText).Methods(http.MethodPost)
+	r.HandleFunc("/eventstreams/multiline", eventstreams.HandleEventStreamMultiLine).Methods(http.MethodPost)
+	r.HandleFunc("/eventstreams/rich", eventstreams.HandleEventStreamRich).Methods(http.MethodPost)
+	r.HandleFunc("/eventstreams/chat", eventstreams.HandleEventStreamChat).Methods(http.MethodPost)
+
+	bind := ":8080"
+	if bindArg != nil {
+		bind = *bindArg
+	}
 
-	log.Println("Listening on :8080")
-	if err := http.ListenAndServe(":8080", r); err != nil {
+	log.Printf("Listening on %s\n", bind)
+	if err := http.ListenAndServe(bind, r); err != nil {
 		log.Fatal(err)
 	}
 }
diff --git a/go.mod b/go.mod
index d6953ae..7b49cce 100644
--- a/go.mod
+++ b/go.mod
@@ -2,4 +2,4 @@ module github.com/speakeasy-api/speakeasy-api-test-service
 
 go 1.19
 
-require github.com/gorilla/mux v1.8.0 // indirect
+require github.com/gorilla/mux v1.8.0
diff --git a/internal/eventstreams/service.go b/internal/eventstreams/service.go
new file mode 100644
index 0000000..d5ae8e8
--- /dev/null
+++ b/internal/eventstreams/service.go
@@ -0,0 +1,128 @@
+package eventstreams
+
+import (
+	"fmt"
+	"net/http"
+	"time"
+)
+
+func pushEvents(rw http.ResponseWriter, events [][]string) {
+	for _, event := range events {
+		for _, line := range event {
+			fmt.Fprintln(rw, line)
+		}
+		fmt.Fprintln(rw, "")
+
+		if f, ok := rw.(http.Flusher); ok {
+			f.Flush()
+		}
+
+		time.Sleep(100 * time.Millisecond)
+	}
+}
+
+func HandleEventStreamJSON(rw http.ResponseWriter, _ *http.Request) {
+	rw.Header().Add("Content-Type", "text/event-stream")
+
+	pushEvents(rw, [][]string{
+		{
+			`data: {"content": "Hello"}`,
+		},
+
+		{
+			`data: {"content": " "}`,
+		},
+
+		{
+			`data: {"content": "world"}`,
+		},
+
+		{
+			`data: {"content": "!"}`,
+		},
+	})
+}
+
+func HandleEventStreamText(rw http.ResponseWriter, _ *http.Request) {
+	rw.Header().Add("Content-Type", "text/event-stream")
+
+	pushEvents(rw, [][]string{
+		{
+			`data: Hello`,
+		},
+
+		{
+			`data:  `,
+		},
+
+		{
+			`data: world`,
+		},
+
+		{
+			`data: !`,
+		},
+	})
+}
+
+func HandleEventStreamMultiLine(rw http.ResponseWriter, _ *http.Request) {
+	rw.Header().Add("Content-Type", "text/event-stream")
+
+	pushEvents(rw, [][]string{
+		{
+			`data: YHOO`,
+			`data: +2`,
+			`data: 10`,
+		},
+	})
+}
+
+func HandleEventStreamRich(rw http.ResponseWriter, _ *http.Request) {
+	rw.Header().Add("Content-Type", "text/event-stream")
+
+	pushEvents(rw, [][]string{
+		{
+			`id: job-1`,
+			`event: completion`,
+			`data: {"completion": "Hello", "stop_reason": null, "model": "jeeves-1"}`,
+		},
+
+		{
+			`event: heartbeat`,
+			`data: ping`,
+			`retry: 3000`,
+		},
+
+		{
+			`id: job-1`,
+			`event: completion`,
+			`data: {"completion": "world!", "stop_reason": "stop_sequence", "model": "jeeves-1"}`,
+		},
+	})
+}
+
+func HandleEventStreamChat(rw http.ResponseWriter, _ *http.Request) {
+	rw.Header().Add("Content-Type", "text/event-stream")
+
+	pushEvents(rw, [][]string{
+		{
+			`data: {"content": "Hello"}`,
+		},
+
+		{
+			`data: {"content": " "}`,
+		},
+
+		{
+			`data: {"content": "world"}`,
+		},
+
+		{
+			`data: {"content": "!"}`,
+		},
+
+		{
+			`data: [DONE]`,
+		},
+	})
+}