Add projects

backend/internal/client/grokbuild/keepalive_test.go

@@ -0,0 +1,156 @@
+package grokbuild
+
+import (
+	"bytes"
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/gin-gonic/gin"
+)
+
+func TestIsGrokClientUserAgent(t *testing.T) {
+	tests := []struct {
+		ua   string
+		want bool
+	}{
+		{"grok-shell/0.2.119 (macos; aarch64)", true},
+		{"grok-pager/1.0.5 grok-shell/1.0.5 (linux; x86_64)", true},
+		{"grok-pager/1.0.5", true},
+		{"GROK-PAGER/1.0", true},
+		{"GROK-SHELL/1.0", true},
+		{"curl/8.7.1", false},
+		{"openai-python/1.0.0", false},
+		{"", false},
+	}
+	for _, tc := range tests {
+		if got := IsGrokClientUserAgent(tc.ua); got != tc.want {
+			t.Errorf("IsGrokClientUserAgent(%q) = %v, want %v", tc.ua, got, tc.want)
+		}
+	}
+}
+
+func TestIsGrokClientHeaders(t *testing.T) {
+	tests := []struct {
+		name    string
+		headers http.Header
+		want    bool
+	}{
+		{
+			name:    "User-Agent with grok-pager",
+			headers: http.Header{"User-Agent": []string{"grok-pager/1.0.5"}},
+			want:    true,
+		},
+		{
+			name:    "case insensitive header name",
+			headers: http.Header{"user-agent": []string{"grok-shell/0.2"}},
+			want:    true,
+		},
+		{
+			name:    "unrelated user agent",
+			headers: http.Header{"User-Agent": []string{"curl/8.7.1"}},
+			want:    false,
+		},
+		{
+			name:    "nil headers",
+			headers: nil,
+			want:    false,
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			if got := IsGrokClientHeaders(tc.headers); got != tc.want {
+				t.Errorf("IsGrokClientHeaders() = %v, want %v", got, tc.want)
+			}
+		})
+	}
+}
+
+func TestIsGrokClientContext(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+	c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
+	c.Request.Header.Set("User-Agent", "grok-pager/1.0.5 grok-shell/1.0.5")
+
+	ctx := context.WithValue(context.Background(), "gin", c)
+	if !IsGrokClientContext(ctx, nil) {
+		t.Error("expected IsGrokClientContext to detect gin context user agent")
+	}
+
+	plainCtx := context.Background()
+	headers := http.Header{"User-Agent": []string{"grok-shell/1.0"}}
+	if !IsGrokClientContext(plainCtx, headers) {
+		t.Error("expected IsGrokClientContext to detect headers when gin context is absent")
+	}
+}
+
+func TestIsKeepalivePayload(t *testing.T) {
+	tests := []struct {
+		payload []byte
+		want    bool
+	}{
+		{[]byte(`{"type":"keepalive","sequence_number":3}`), true},
+		{[]byte(`{"type":"keepalive"}`), true},
+		{[]byte(`{"type":"response.created"}`), false},
+		{[]byte(`{"type":"response.reasoning.delta"}`), false},
+		{[]byte(``), false},
+	}
+	for _, tc := range tests {
+		if got := IsKeepalivePayload(tc.payload); got != tc.want {
+			t.Errorf("IsKeepalivePayload(%s) = %v, want %v", string(tc.payload), got, tc.want)
+		}
+	}
+}
+
+func TestIsKeepaliveSSELine(t *testing.T) {
+	tests := []struct {
+		line []byte
+		want bool
+	}{
+		{[]byte("event: keepalive"), true},
+		{[]byte("event: keepalive\n"), true},
+		{[]byte("  event: keepalive  "), true},
+		{[]byte(`data: {"type":"keepalive","sequence_number":3}`), true},
+		{[]byte(`data: {"type":"keepalive"}`), true},
+		{[]byte("event: response.created"), false},
+		{[]byte("event: keepalive-other"), false},
+		{[]byte(`data: {"type":"response.created"}`), false},
+		{[]byte(""), false},
+	}
+	for _, tc := range tests {
+		if got := IsKeepaliveSSELine(tc.line); got != tc.want {
+			t.Errorf("IsKeepaliveSSELine(%s) = %v, want %v", string(tc.line), got, tc.want)
+		}
+	}
+}
+
+func TestTransformKeepaliveSSELine(t *testing.T) {
+	comment := KeepaliveSSEComment()
+
+	// Grok client: keepalive line is transformed
+	got, ok := TransformKeepaliveSSELine([]byte("event: keepalive"), true)
+	if !ok || !bytes.Equal(got, comment) {
+		t.Errorf("TransformKeepaliveSSELine(event: keepalive, true) = %q, %v, want %q, true", string(got), ok, string(comment))
+	}
+
+	got, ok = TransformKeepaliveSSELine([]byte(`data: {"type":"keepalive","sequence_number":3}`), true)
+	if !ok || !bytes.Equal(got, comment) {
+		t.Errorf("TransformKeepaliveSSELine(data: keepalive, true) = %q, %v, want %q, true", string(got), ok, string(comment))
+	}
+
+	// Grok client: normal line is untouched
+	normalLine := []byte(`data: {"type":"response.created"}`)
+	got, ok = TransformKeepaliveSSELine(normalLine, true)
+	if ok || !bytes.Equal(got, normalLine) {
+		t.Errorf("TransformKeepaliveSSELine(normalLine, true) = %q, %v, want unchanged, false", string(got), ok)
+	}
+
+	// Non-Grok client: keepalive line is untouched
+	keepaliveLine := []byte("event: keepalive")
+	got, ok = TransformKeepaliveSSELine(keepaliveLine, false)
+	if ok || !bytes.Equal(got, keepaliveLine) {
+		t.Errorf("TransformKeepaliveSSELine(event: keepalive, false) = %q, %v, want unchanged, false", string(got), ok)
+	}
+}