vibe-proxy/backend/internal/client/grokbuild/keepalive_test.go

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)
	}
}