api key integration

services/assistantLimit_test.go

@@ -0,0 +1,113 @@
+package services
+
+import (
+	"errors"
+	"testing"
+	"time"
+)
+
+// The rate limit, which is a cost control and not a security one.
+
+func TestABurstOfQuestionsGetsThrough(t *testing.T) {
+	// Somebody working through the prompt chips on a page fires four in a row
+	// and must not be stopped mid-thought.
+	at := time.Now()
+	l := newAskLimiter(func() time.Time { return at })
+
+	for i := range askBurst {
+		if err := l.allow(904); err != nil {
+			t.Fatalf("question %d of a burst was refused: %v", i+1, err)
+		}
+	}
+}
+
+func TestTheSeventhQuestionInARowWaits(t *testing.T) {
+	at := time.Now()
+	l := newAskLimiter(func() time.Time { return at })
+	for range askBurst {
+		_ = l.allow(904)
+	}
+
+	err := l.allow(904)
+	var tooFast ErrTooFast
+	if !errors.As(err, &tooFast) {
+		t.Fatalf("an unbounded burst was allowed: %v", err)
+	}
+	if tooFast.RetryIn <= 0 {
+		t.Fatal("the refusal does not say how long to wait")
+	}
+	// The message is read by a shopkeeper, so it has to say something useful.
+	if !contains(err.Error(), "seconds") {
+		t.Fatalf("unhelpful message: %q", err)
+	}
+}
+
+func TestTheAllowanceComesBackWithTime(t *testing.T) {
+	// Continuous refill, not a fixed window: a person is never told to wait out
+	// a window that started before they arrived.
+	at := time.Now()
+	l := newAskLimiter(func() time.Time { return at })
+	for range askBurst {
+		_ = l.allow(904)
+	}
+	if err := l.allow(904); err == nil {
+		t.Fatal("expected to be out of allowance")
+	}
+
+	at = at.Add(askRefill)
+	if err := l.allow(904); err != nil {
+		t.Fatalf("one refill period bought nothing: %v", err)
+	}
+}
+
+func TestOnePersonCannotLockOutTheirColleagues(t *testing.T) {
+	// Per user rather than per tenant, so an impatient person in a shop cannot
+	// turn a cost control into an outage for everybody else.
+	at := time.Now()
+	l := newAskLimiter(func() time.Time { return at })
+	for range askBurst + 3 {
+		_ = l.allow(904)
+	}
+
+	if err := l.allow(905); err != nil {
+		t.Fatalf("a colleague was refused because somebody else was busy: %v", err)
+	}
+}
+
+func TestAnUnidentifiedCallerIsNotSilentlyThrottled(t *testing.T) {
+	// Every such request is already refused for having no session. Rate
+	// limiting it too would hide the real reason behind a confusing one.
+	l := newAskLimiter(nil)
+	for range askBurst * 3 {
+		if err := l.allow(0); err != nil {
+			t.Fatalf("a caller with no user id was throttled: %v", err)
+		}
+	}
+}
+
+func TestQuietCallersAreForgotten(t *testing.T) {
+	// So the map does not grow with every account that ever asked anything.
+	at := time.Now()
+	l := newAskLimiter(func() time.Time { return at })
+
+	for user := 1; user <= 300; user++ {
+		_ = l.allow(user)
+	}
+	before := len(l.buckets)
+
+	at = at.Add(askIdle + time.Minute)
+	_ = l.allow(99999) // a new caller is what triggers the sweep
+
+	if len(l.buckets) >= before {
+		t.Fatalf("nothing was swept: %d then %d", before, len(l.buckets))
+	}
+}
+
+func contains(haystack, needle string) bool {
+	for i := 0; i+len(needle) <= len(haystack); i++ {
+		if haystack[i:i+len(needle)] == needle {
+			return true
+		}
+	}
+	return false
+}