review-bot/llm/aicore_test.go

package llm

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

func TestAICoreClient_TokenFetch(t *testing.T) {
	tokenCalls := int32(0)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/oauth/token" {
			atomic.AddInt32(&tokenCalls, 1)
			if r.Method != http.MethodPost {
				t.Errorf("expected POST for token, got %s", r.Method)
			}
			if r.Header.Get("Content-Type") != "application/x-www-form-urlencoded" {
				t.Errorf("expected form content type")
			}
			w.Header().Set("Content-Type", "application/json")
			json.NewEncoder(w).Encode(map[string]interface{}{
				"access_token": "test-token-123",
				"expires_in":   3600,
			})
			return
		}
		t.Errorf("unexpected path: %s", r.URL.Path)
	}))
	defer server.Close()

	client := NewAICoreClient(AICoreConfig{
		ClientID:      "test-id",
		ClientSecret:  "test-secret",
		AuthURL:       server.URL,
		APIURL:        server.URL,
		ResourceGroup: "default",
	})

	token, err := client.getToken(context.Background())
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if token != "test-token-123" {
		t.Errorf("expected token 'test-token-123', got %q", token)
	}

	// Second call should use cached token
	token2, err := client.getToken(context.Background())
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if token2 != "test-token-123" {
		t.Errorf("expected cached token")
	}
	if atomic.LoadInt32(&tokenCalls) != 1 {
		t.Errorf("expected 1 token call (cached), got %d", tokenCalls)
	}
}

func TestAICoreClient_DeploymentFetch(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/oauth/token" {
			w.Header().Set("Content-Type", "application/json")
			json.NewEncoder(w).Encode(map[string]interface{}{
				"access_token": "test-token",
				"expires_in":   3600,
			})
			return
		}
		if r.URL.Path == "/v2/lm/deployments" {
			if r.Header.Get("Authorization") != "Bearer test-token" {
				t.Errorf("expected Bearer auth")
			}
			if r.Header.Get("AI-Resource-Group") != "default" {
				t.Errorf("expected resource group header")
			}
			w.Header().Set("Content-Type", "application/json")
			json.NewEncoder(w).Encode(map[string]interface{}{
				"resources": []map[string]interface{}{
					{
						"id":            "deploy-123",
						"deploymentUrl": "https://example.com/v2/inference/deployments/deploy-123",
						"status":        "RUNNING",
						"details": map[string]interface{}{
							"resources": map[string]interface{}{
								"backend_details": map[string]interface{}{
									"model": map[string]interface{}{
										"name": "anthropic--claude-4.6-sonnet",
									},
								},
							},
						},
					},
					{
						"id":            "deploy-456",
						"deploymentUrl": "https://example.com/v2/inference/deployments/deploy-456",
						"status":        "STOPPED",
						"details": map[string]interface{}{
							"resources": map[string]interface{}{
								"backend_details": map[string]interface{}{
									"model": map[string]interface{}{
										"name": "gpt-5",
									},
								},
							},
						},
					},
					{
						"id":            "deploy-789",
						"deploymentUrl": "https://example.com/v2/inference/deployments/deploy-789",
						"status":        "RUNNING",
						"details": map[string]interface{}{
							"resources": map[string]interface{}{
								"backend_details": map[string]interface{}{
									"model": map[string]interface{}{
										"name": "gpt-5",
									},
								},
							},
						},
					},
				},
			})
			return
		}
		t.Errorf("unexpected path: %s", r.URL.Path)
	}))
	defer server.Close()

	client := NewAICoreClient(AICoreConfig{
		ClientID:      "test-id",
		ClientSecret:  "test-secret",
		AuthURL:       server.URL,
		APIURL:        server.URL,
		ResourceGroup: "default",
	})

	// Should find running deployment
	url, err := client.getDeploymentURL(context.Background(), "anthropic--claude-4.6-sonnet")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if url != "https://example.com/v2/inference/deployments/deploy-123" {
		t.Errorf("unexpected URL: %s", url)
	}

	// Should find running gpt-5, not stopped one
	url, err = client.getDeploymentURL(context.Background(), "gpt-5")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if url != "https://example.com/v2/inference/deployments/deploy-789" {
		t.Errorf("unexpected URL: %s", url)
	}

	// Should error on unknown model
	_, err = client.getDeploymentURL(context.Background(), "unknown-model")
	if err == nil {
		t.Error("expected error for unknown model")
	}
}

func TestAICoreClient_CompleteAnthropic(t *testing.T) {
	// Use a pointer to capture the server URL for use in the handler
	var baseURL string
	mux := http.NewServeMux()
	mux.HandleFunc("/oauth/token", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]interface{}{
			"access_token": "test-token",
			"expires_in":   3600,
		})
	})
	mux.HandleFunc("/v2/lm/deployments", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]interface{}{
			"resources": []map[string]interface{}{
				{
					"id":            "deploy-anthropic",
					"deploymentUrl": baseURL + "/deployments/anthropic",
					"status":        "RUNNING",
					"details": map[string]interface{}{
						"resources": map[string]interface{}{
							"backend_details": map[string]interface{}{
								"model": map[string]interface{}{
									"name": "anthropic--claude-4.6-sonnet",
								},
							},
						},
					},
				},
			},
		})
	})
	mux.HandleFunc("/deployments/anthropic/invoke", func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer test-token" {
			t.Errorf("expected Bearer auth on invoke")
		}
		var req anthropicRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			t.Fatalf("decode request: %v", err)
		}
		if req.Model != "anthropic--claude-4.6-sonnet" {
			t.Errorf("expected model name in request")
		}
		if req.System != "You are helpful" {
			t.Errorf("expected system prompt: %q", req.System)
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]interface{}{
			"content": []map[string]interface{}{
				{"type": "text", "text": "Hello from AI Core!"},
			},
		})
	})

	server := httptest.NewServer(mux)
	baseURL = server.URL
	defer server.Close()

	client := NewAICoreClient(AICoreConfig{
		ClientID:      "test-id",
		ClientSecret:  "test-secret",
		AuthURL:       server.URL,
		APIURL:        server.URL,
		ResourceGroup: "default",
	})

	result, err := client.CompleteAnthropic(context.Background(), "anthropic--claude-4.6-sonnet", []Message{
		{Role: "system", Content: "You are helpful"},
		{Role: "user", Content: "Hello"},
	}, 8192, 0)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if result != "Hello from AI Core!" {
		t.Errorf("expected 'Hello from AI Core!', got %q", result)
	}
}

func TestAICoreClient_CompleteOpenAI(t *testing.T) {
	var baseURL string
	mux := http.NewServeMux()
	mux.HandleFunc("/oauth/token", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]interface{}{
			"access_token": "test-token",
			"expires_in":   3600,
		})
	})
	mux.HandleFunc("/v2/lm/deployments", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]interface{}{
			"resources": []map[string]interface{}{
				{
					"id":            "deploy-openai",
					"deploymentUrl": baseURL + "/deployments/openai",
					"status":        "RUNNING",
					"details": map[string]interface{}{
						"resources": map[string]interface{}{
							"backend_details": map[string]interface{}{
								"model": map[string]interface{}{
									"name": "gpt-5",
								},
							},
						},
					},
				},
			},
		})
	})
	mux.HandleFunc("/deployments/openai/chat/completions", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Query().Get("api-version") != "2024-12-01-preview" {
			t.Errorf("expected api-version query param")
		}
		var req ChatRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			t.Fatalf("decode request: %v", err)
		}
		if req.Model != "gpt-5" {
			t.Errorf("expected model gpt-5, got %s", req.Model)
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(ChatResponse{
			Choices: []struct {
				Message struct {
					Content string `json:"content"`
				} `json:"message"`
			}{
				{Message: struct {
					Content string `json:"content"`
				}{Content: "Hello from GPT-5!"}},
			},
		})
	})

	server := httptest.NewServer(mux)
	baseURL = server.URL
	defer server.Close()

	client := NewAICoreClient(AICoreConfig{
		ClientID:      "test-id",
		ClientSecret:  "test-secret",
		AuthURL:       server.URL,
		APIURL:        server.URL,
		ResourceGroup: "default",
	})

	result, err := client.CompleteOpenAI(context.Background(), "gpt-5", []Message{
		{Role: "user", Content: "Hello"},
	}, 0)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if result != "Hello from GPT-5!" {
		t.Errorf("expected 'Hello from GPT-5!', got %q", result)
	}
}

func TestIsAnthropicModel(t *testing.T) {
	tests := []struct {
		model    string
		expected bool
	}{
		{"anthropic--claude-4.6-sonnet", true},
		{"anthropic--claude-4.6-opus", true},
		{"claude-sonnet-4", true},
		{"gpt-5", false},
		{"gpt-4.1", false},
		{"llama-3", false},
	}

	for _, tt := range tests {
		got := IsAnthropicModel(tt.model)
		if got != tt.expected {
			t.Errorf("IsAnthropicModel(%q) = %v, want %v", tt.model, got, tt.expected)
		}
	}
}

func TestAICoreClient_TokenExpiry(t *testing.T) {
	tokenCalls := int32(0)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/oauth/token" {
			call := atomic.AddInt32(&tokenCalls, 1)
			w.Header().Set("Content-Type", "application/json")
			json.NewEncoder(w).Encode(map[string]interface{}{
				"access_token": "token-" + string(rune('0'+call)),
				"expires_in":   1, // 1 second expiry
			})
			return
		}
	}))
	defer server.Close()

	client := NewAICoreClient(AICoreConfig{
		ClientID:      "test-id",
		ClientSecret:  "test-secret",
		AuthURL:       server.URL,
		APIURL:        server.URL,
		ResourceGroup: "default",
	})

	// First call
	token1, _ := client.getToken(context.Background())

	// Force token expiry by manipulating expiry time
	client.mu.Lock()
	client.tokenExpiry = time.Now().Add(-time.Hour)
	client.mu.Unlock()

	// Should fetch new token
	token2, _ := client.getToken(context.Background())

	if token1 == token2 {
		t.Error("expected different tokens after expiry")
	}
	if atomic.LoadInt32(&tokenCalls) != 2 {
		t.Errorf("expected 2 token calls, got %d", tokenCalls)
	}
}

func TestClient_WithAICore(t *testing.T) {
	client := NewClient("http://example.com", "key", "model")
	if client.provider != ProviderOpenAI {
		t.Errorf("expected default provider openai, got %s", client.provider)
	}

	client.WithAICore(AICoreConfig{
		ClientID:      "id",
		ClientSecret:  "secret",
		AuthURL:       "https://auth.example.com",
		APIURL:        "https://api.example.com",
		ResourceGroup: "default",
	})

	if client.provider != ProviderAICore {
		t.Errorf("expected provider aicore, got %s", client.provider)
	}
	if client.aicore == nil {
		t.Error("expected aicore client to be set")
	}
}

func TestClient_CompleteAICore(t *testing.T) {
	var baseURL string
	mux := http.NewServeMux()
	mux.HandleFunc("/oauth/token", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]interface{}{
			"access_token": "test-token",
			"expires_in":   3600,
		})
	})
	mux.HandleFunc("/v2/lm/deployments", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]interface{}{
			"resources": []map[string]interface{}{
				{
					"id":            "deploy-test",
					"deploymentUrl": baseURL + "/deployments/test",
					"status":        "RUNNING",
					"details": map[string]interface{}{
						"resources": map[string]interface{}{
							"backend_details": map[string]interface{}{
								"model": map[string]interface{}{
									"name": "gpt-5",
								},
							},
						},
					},
				},
			},
		})
	})
	mux.HandleFunc("/deployments/test/chat/completions", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(ChatResponse{
			Choices: []struct {
				Message struct {
					Content string `json:"content"`
				} `json:"message"`
			}{
				{Message: struct {
					Content string `json:"content"`
				}{Content: "AI Core via Client works!"}},
			},
		})
	})

	server := httptest.NewServer(mux)
	baseURL = server.URL
	defer server.Close()

	client := NewClient("", "", "gpt-5").WithAICore(AICoreConfig{
		ClientID:      "test-id",
		ClientSecret:  "test-secret",
		AuthURL:       server.URL,
		APIURL:        server.URL,
		ResourceGroup: "default",
	})

	result, err := client.Complete(context.Background(), []Message{
		{Role: "user", Content: "Hello"},
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if !strings.Contains(result, "AI Core via Client works!") {
		t.Errorf("unexpected result: %s", result)
	}
}