integration tests even if they are slow

vicuna-tools.d/main.go

@@ -90,7 +90,7 @@ func config(ctx context.Context) {
 	flag.StringVar(&Config.ChatBot.SessionD, "chatbot-session-d", d, "dir to store chat bot sessions")
 	flag.StringVar(&Config.ChatBot.WD, "chatbot-working-d", "./llama.cpp", "working directory for chatbot")
 	flag.StringVar(&Config.ChatBot.Command, "chatbot-cmd", "./main -m ./models/ggml-vic7b-uncensored-q5_1.bin --repeat_penalty 1.0", "chatbot cmd prefix")
-	flag.IntVar(&Config.ChatBot.N, "chatbot-n", 64, "chatbot items to gen")
+	flag.IntVar(&Config.ChatBot.N, "chatbot-n", 256, "chatbot items to gen")
 	flag.BoolVar(&Config.Debug, "debug", false, "debug mode")
 	flag.Parse()
 }
@@ -149,7 +149,7 @@ func handleLogin(w http.ResponseWriter, r *http.Request) error {
 		if err != nil {
 			return err
 		}
-		cookie.Serialize(w)
+		cookie.Inject(w)
 		http.Redirect(w, r, "/", http.StatusTemporaryRedirect)
 		return nil
 	default:
@@ -215,16 +215,19 @@ func (cookie Cookie) Verify() error {
 	return nil
 }
 
-func (cookie Cookie) Serialize(w http.ResponseWriter) {
-	b, _ := json.Marshal(cookie)
-	encoded := base64.URLEncoding.EncodeToString(b)
+func (cookie Cookie) Inject(w http.ResponseWriter) {
 	c := &http.Cookie{
 		Name:  "root",
-		Value: encoded,
+		Value: cookie.Serialize(),
 	}
 	http.SetCookie(w, c)
 }
 
+func (cookie Cookie) Serialize() string {
+	b, _ := json.Marshal(cookie)
+	return base64.URLEncoding.EncodeToString(b)
+}
+
 func handleAPI(w http.ResponseWriter, r *http.Request) error {
 	if _, err := ParseCookie(r); err != nil {
 		http.Redirect(w, r, "/login", http.StatusTemporaryRedirect)
@@ -344,7 +347,7 @@ func handleAPIChatBotPut(w http.ResponseWriter, r *http.Request) error {
 		commands[1:]...,
 	)
 	command.Dir = Config.ChatBot.WD
-	command.Stderr = os.Stderr
+	command.Stderr = log.Writer()
 
 	stdout, err := command.StdoutPipe()
 	if err != nil {
@@ -400,5 +403,6 @@ func appendFile(toF, msg string) error {
 	if _, err := f.WriteString(msg); err != nil {
 		return err
 	}
+	f.Write([]byte("\n"))
 	return nil
 }

vicuna-tools.d/main_integration_test.go

@@ -0,0 +1,122 @@
+//go:build integration
+
+package main
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"net/url"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestAPIV0ChatBot(t *testing.T) {
+	defer goTestMain(t)()
+
+	body := url.Values{}
+	body.Set(`Prompt`, "Text transcript of a never ending dialogue between a middle manager and his direct reports. The middle manager works in a middle sized corporation and must tell employees what he thinks of employees' work. Middle manager always replies to bad news with an overly optimistic observation prefixed with 'Perhaps, but have you considered'.")
+	body.Set(`Message`, `I lost keys to the company car in my couch, boss.`)
+
+	t.Run("put over post", func(t *testing.T) {
+		resp := httpDo(t, http.MethodPost, "/api/v0/chatbot", body.Encode())
+		got, err := io.ReadAll(resp.Body)
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Logf("(%d) %s", resp.StatusCode, got)
+
+		resp2 := httpDo(t, http.MethodPut, "/api/v0/chatbot", body.Encode())
+		got2, err := io.ReadAll(resp2.Body)
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Logf("(%d) %s", resp2.StatusCode, got2)
+
+		if len(got) == len(got2) {
+			t.Error("nothing new as of put")
+		} else if !bytes.HasPrefix(got2, got) {
+			t.Error("put was not a continuation")
+		}
+	})
+
+	t.Run("post over post", func(t *testing.T) {
+		resp := httpDo(t, http.MethodPost, "/api/v0/chatbot", body.Encode())
+		got, err := io.ReadAll(resp.Body)
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Logf("(%d) %s", resp.StatusCode, got)
+
+		resp2 := httpDo(t, http.MethodPost, "/api/v0/chatbot", body.Encode())
+		got2, err := io.ReadAll(resp2.Body)
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Logf("(%d) %s", resp2.StatusCode, got2)
+
+		if bytes.HasPrefix(got2, got) {
+			t.Error("post over post was a continuation")
+		}
+	})
+
+	t.Run("put over zero", func(t *testing.T) {
+		resp := httpDo(t, http.MethodPut, "/api/v0/chatbot", body.Encode())
+		got, err := io.ReadAll(resp.Body)
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Logf("(%d) %s", resp.StatusCode, got)
+	})
+
+}
+
+func goTestMain(t *testing.T) func() {
+	log.SetOutput(io.Discard)
+
+	ctx, can := context.WithCancel(context.Background())
+	ctx, cleanup := contextWithCleanup(ctx)
+	config(ctx)
+	Config.Port += 10
+	wg := &sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		listenAndServe(ctx)
+	}()
+	httpDo(t, http.MethodGet, "/", "")
+	return func() {
+		cleanup()
+		can()
+		wg.Wait()
+	}
+}
+
+func httpDo(t *testing.T, method, path, body string) *http.Response {
+	req, _ := http.NewRequest(
+		method,
+		fmt.Sprintf("http://localhost:%d/%s", Config.Port, strings.TrimLeft(path, "/")),
+		nil,
+	)
+	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+	cookie := Cookie{Name: t.Name()}
+	req.Header.Set("Cookie", "root="+cookie.Serialize())
+	for {
+		req.Body = io.NopCloser(strings.NewReader(body))
+		if resp, err := http.DefaultClient.Do(req.Clone(context.Background())); err == nil {
+			defer resp.Body.Close()
+			b, _ := io.ReadAll(resp.Body)
+			resp.Body = io.NopCloser(bytes.NewReader(b))
+			return resp
+		}
+		time.Sleep(time.Millisecond * 25)
+		t.Logf("retrying %s", req.URL.String())
+	}
+	t.Fatalf("failed to ever %s %s", method, path)
+	panic(nil)
+}