iris-context/tests/ai/client_test.go

package tests

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/versenilvis/iris/ai"
	"github.com/versenilvis/iris/config"
)

func TestCleanSuggestion(t *testing.T) {
	tests := []struct {
		input    string
		expected string
	}{
		{"  docker run -d nginx  ", "docker run -d nginx"},
		{"```bash\ngit status\n```", "git status"},
		{"```\nls -la\n```", "ls -la"},
		{"`npm run dev`", "npm run dev"},
		{"\"docker ps\"", "docker ps"},
		{"'git diff'", "git diff"},
		{"\"git commit -m 'hello'\"", "\"git commit -m 'hello'\""},
	}

	for _, tt := range tests {
		got := ai.CleanSuggestion(tt.input)
		if got != tt.expected {
			t.Errorf("CleanSuggestion(%q) = %q, want %q", tt.input, got, tt.expected)
		}
	}
}

func TestOpenAIClient_Suggest(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			t.Errorf("expected post method, got %s", r.Method)
		}
		if r.Header.Get("Authorization") != "Bearer test-secret-key" {
			t.Errorf("expected bearer token, got %s", r.Header.Get("Authorization"))
		}
		if r.Header.Get("Content-Type") != "application/json" {
			t.Errorf("expected application/json, got %s", r.Header.Get("Content-Type"))
		}

		body, err := io.ReadAll(r.Body)
		if err != nil {
			t.Fatalf("failed to read request body: %v", err)
		}

		var reqMap map[string]any
		if err := json.Unmarshal(body, &reqMap); err != nil {
			t.Fatalf("failed to parse request json: %v", err)
		}
		if reqMap["model"] != "test-model-32b" {
			t.Errorf("expected model test-model-32b, got %v", reqMap["model"])
		}
		if reqMap["temperature"] != 0.5 {
			t.Errorf("expected extra temperature 0.5, got %v", reqMap["temperature"])
		}

		res := map[string]any{
			"choices": []map[string]any{
				{"message": map[string]any{"role": "assistant", "content": "```bash\nkubectl get pods -n kube-system\n```"}},
			},
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(res)
	}))
	defer server.Close()

	cfg := config.ProviderConfig{
		InheritedFrom: "openai",
		Endpoint:      server.URL,
		APIKey:        "test-secret-key",
		Model:         "test-model-32b",
		TimeoutMS:     1000,
		ExtraRequestBody: map[string]any{
			"temperature": 0.5,
		},
	}

	client, err := ai.NewClient(cfg)
	if err != nil {
		t.Fatalf("failed to create client: %v", err)
	}

	ctx := context.Background()
	env := ai.EnvSnapshot{Cwd: "/home/user", LastCmd: "kubectl get", LastExitCode: 0}
	sugg, err := client.Suggest(ctx, "kubectl get p", env, "")
	if err != nil {
		t.Fatalf("suggest failed: %v", err)
	}
	if sugg == nil {
		t.Fatalf("expected suggestion, got nil")
	}
	if sugg.Cmd != "kubectl get pods -n kube-system" {
		t.Errorf("expected cleaned cmd, got %q", sugg.Cmd)
	}
	if sugg.Confidence != 85 {
		t.Errorf("expected confidence 85, got %d", sugg.Confidence)
	}
}

func TestOpenAIClient_TimeoutAndCancel(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(200 * time.Millisecond)
		w.WriteHeader(http.StatusOK)
	}))
	defer server.Close()

	cfg := config.ProviderConfig{
		InheritedFrom: "openai",
		Endpoint:      server.URL,
		TimeoutMS:     50,
	}

	client := ai.NewOpenAIClient(cfg)
	ctx := context.Background()
	env := ai.EnvSnapshot{}

	_, err := client.Suggest(ctx, "sleep", env, "")
	if err == nil {
		t.Errorf("expected timeout error, got nil")
	}

	ctxCancel, cancel := context.WithCancel(context.Background())
	cancel()
	_, err = client.Suggest(ctxCancel, "sleep", env, "")
	if err == nil {
		t.Errorf("expected context canceled error, got nil")
	}
}