plugins/mocker/plugin_test.go

package mocker

import (
	"context"
	"testing"

	bifrost "github.com/maximhq/bifrost/core"
	"github.com/maximhq/bifrost/core/schemas"
)

// BaseAccount implements the schemas.Account interface for testing purposes.
// It provides mock implementations of the required methods to test the Mocker plugin
// with a basic OpenAI configuration.
type BaseAccount struct{}

// GetConfiguredProviders returns a list of supported providers for testing.
func (baseAccount *BaseAccount) GetConfiguredProviders() ([]schemas.ModelProvider, error) {
	return []schemas.ModelProvider{schemas.OpenAI, schemas.Anthropic}, nil
}

// GetKeysForProvider returns a dummy API key configuration for testing.
// Since we're testing the mocker plugin, these keys should never be used
// as the plugin intercepts requests before they reach the actual providers.
func (baseAccount *BaseAccount) GetKeysForProvider(ctx context.Context, providerKey schemas.ModelProvider) ([]schemas.Key, error) {
	return []schemas.Key{
		{
			Value:  *schemas.NewEnvVar("dummy-api-key-for-testing"), // Dummy key
			Models: []string{"gpt-4", "gpt-4-turbo", "claude-3"},
			Weight: 1.0,
		},
	}, nil
}

// GetConfigForProvider returns default provider configuration for testing.
func (baseAccount *BaseAccount) GetConfigForProvider(providerKey schemas.ModelProvider) (*schemas.ProviderConfig, error) {
	return &schemas.ProviderConfig{
		NetworkConfig:            schemas.DefaultNetworkConfig,
		ConcurrencyAndBufferSize: schemas.DefaultConcurrencyAndBufferSize,
	}, nil
}

func (baseAccount *BaseAccount) GetGlobalKeys(ctx context.Context) ([]schemas.Key, error) {
	return nil, nil
}

// TestMockerPlugin_GetName tests the plugin name
func TestMockerPlugin_GetName(t *testing.T) {
	plugin, err := Init(MockerConfig{})
	if err != nil {
		t.Fatalf("Expected no error creating plugin, got: %v", err)
	}
	if plugin.GetName() != PluginName {
		t.Errorf("Expected '%s', got '%s'", PluginName, plugin.GetName())
	}
}

// TestMockerPlugin_Disabled tests that disabled plugin doesn't interfere
func TestMockerPlugin_Disabled(t *testing.T) {
	ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
	config := MockerConfig{
		Enabled: false,
	}
	plugin, err := Init(config)
	if err != nil {
		t.Fatalf("Expected no error creating plugin, got: %v", err)
	}

	account := BaseAccount{}
	client, err := bifrost.Init(ctx, schemas.BifrostConfig{
		Account:    &account,
		LLMPlugins: []schemas.LLMPlugin{plugin},
		Logger:     bifrost.NewDefaultLogger(schemas.LogLevelError),
	})
	if err != nil {
		t.Fatalf("Error initializing Bifrost: %v", err)
	}
	defer client.Shutdown()

	// This should pass through to the real provider (but will fail due to dummy key)
	_, bifrostErr := client.ChatCompletionRequest(ctx, &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: bifrost.Ptr("Hello, test message"),
				},
			},
		},
	})

	// Should get an authentication error from OpenAI, not a mock response
	// This proves the plugin is disabled and not intercepting requests
	if bifrostErr == nil {
		t.Error("Expected error from real provider with dummy API key")
	}
}

// TestMockerPlugin_DefaultMockRule tests the default catch-all rule
func TestMockerPlugin_DefaultMockRule(t *testing.T) {
	ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
	config := MockerConfig{
		Enabled: true, // No rules provided, should create default rule
	}
	plugin, err := Init(config)
	if err != nil {
		t.Fatalf("Expected no error creating plugin, got: %v", err)
	}

	account := BaseAccount{}
	client, err := bifrost.Init(ctx, schemas.BifrostConfig{
		Account:    &account,
		LLMPlugins: []schemas.LLMPlugin{plugin},
		Logger:     bifrost.NewDefaultLogger(schemas.LogLevelError),
	})
	if err != nil {
		t.Fatalf("Error initializing Bifrost: %v", err)
	}
	defer client.Shutdown()

	response, bifrostErr := client.ChatCompletionRequest(ctx, &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: bifrost.Ptr("Hello, test message"),
				},
			},
		},
	})

	if bifrostErr != nil {
		t.Fatalf("Expected no error, got: %v", bifrostErr)
	}
	if response == nil {
		t.Fatal("Expected response")
	}
	if len(response.Choices) == 0 {
		t.Fatal("Expected at least one choice")
	}
	if response.Choices[0].ChatNonStreamResponseChoice.Message.Content.ContentStr == nil {
		t.Fatal("Expected content string")
	}
	if *response.Choices[0].ChatNonStreamResponseChoice.Message.Content.ContentStr != "This is a mock response from the Mocker plugin" {
		t.Errorf("Expected default mock message, got: %s", *response.Choices[0].ChatNonStreamResponseChoice.Message.Content.ContentStr)
	}
}

// TestMockerPlugin_CustomSuccessRule tests custom success response
func TestMockerPlugin_CustomSuccessRule(t *testing.T) {
	ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
	config := MockerConfig{
		Enabled: true,
		Rules: []MockRule{
			{
				Name:        "openai-success",
				Enabled:     true,
				Priority:    100,
				Probability: 1.0,
				Conditions: Conditions{
					Providers: []string{"openai"},
				},
				Responses: []Response{
					{
						Type: ResponseTypeSuccess,
						Content: &SuccessResponse{
							Message: "Custom OpenAI mock response",
							Usage: &Usage{
								PromptTokens:     15,
								CompletionTokens: 25,
								TotalTokens:      40,
							},
						},
					},
				},
			},
		},
	}
	plugin, err := Init(config)
	if err != nil {
		t.Fatalf("Expected no error creating plugin, got: %v", err)
	}

	account := BaseAccount{}
	client, err := bifrost.Init(ctx, schemas.BifrostConfig{
		Account:    &account,
		LLMPlugins: []schemas.LLMPlugin{plugin},
		Logger:     bifrost.NewDefaultLogger(schemas.LogLevelError),
	})
	if err != nil {
		t.Fatalf("Error initializing Bifrost: %v", err)
	}
	defer client.Shutdown()

	response, bifrostErr := client.ChatCompletionRequest(ctx, &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: bifrost.Ptr("Hello, test message"),
				},
			},
		},
	})

	if bifrostErr != nil {
		t.Fatalf("Expected no error, got: %v", bifrostErr)
	}
	if response == nil {
		t.Fatal("Expected response")
	}
	if len(response.Choices) == 0 {
		t.Fatal("Expected at least one choice")
	}
	if response.Choices[0].ChatNonStreamResponseChoice.Message.Content.ContentStr == nil {
		t.Fatal("Expected content string")
	}
	if *response.Choices[0].ChatNonStreamResponseChoice.Message.Content.ContentStr != "Custom OpenAI mock response" {
		t.Errorf("Expected custom message, got: %s", *response.Choices[0].ChatNonStreamResponseChoice.Message.Content.ContentStr)
	}
	if response.Usage.TotalTokens != 40 {
		t.Errorf("Expected 40 total tokens, got %d", response.Usage.TotalTokens)
	}
}

// TestMockerPlugin_ErrorResponse tests error response generation
func TestMockerPlugin_ErrorResponse(t *testing.T) {
	ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
	allowFallbacks := false
	config := MockerConfig{
		Enabled: true,
		Rules: []MockRule{
			{
				Name:        "rate-limit-error",
				Enabled:     true,
				Priority:    100,
				Probability: 1.0,
				Conditions: Conditions{
					Providers: []string{"openai"},
				},
				Responses: []Response{
					{
						Type:           ResponseTypeError,
						AllowFallbacks: &allowFallbacks,
						Error: &ErrorResponse{
							Message:    "Rate limit exceeded",
							Type:       bifrost.Ptr("rate_limit"),
							Code:       bifrost.Ptr("429"),
							StatusCode: bifrost.Ptr(429),
						},
					},
				},
			},
		},
	}
	plugin, err := Init(config)
	if err != nil {
		t.Fatalf("Expected no error creating plugin, got: %v", err)
	}

	account := BaseAccount{}
	client, err := bifrost.Init(ctx, schemas.BifrostConfig{
		Account:    &account,
		LLMPlugins: []schemas.LLMPlugin{plugin},
		Logger:     bifrost.NewDefaultLogger(schemas.LogLevelError),
	})
	if err != nil {
		t.Fatalf("Error initializing Bifrost: %v", err)
	}
	defer client.Shutdown()

	_, bifrostErr := client.ChatCompletionRequest(ctx, &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: bifrost.Ptr("Hello, test message"),
				},
			},
		},
	})

	if bifrostErr == nil {
		t.Fatal("Expected error response")
	}
	if bifrostErr.Error.Message != "Rate limit exceeded" {
		t.Errorf("Expected 'Rate limit exceeded', got: %s", bifrostErr.Error.Message)
	}
	if bifrostErr.StatusCode == nil || *bifrostErr.StatusCode != 429 {
		t.Errorf("Expected status code 429, got: %v", bifrostErr.StatusCode)
	}
}

// TestMockerPlugin_MessageTemplate tests template variable substitution
func TestMockerPlugin_MessageTemplate(t *testing.T) {
	ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
	config := MockerConfig{
		Enabled: true,
		Rules: []MockRule{
			{
				Name:        "template-test",
				Enabled:     true,
				Priority:    100,
				Probability: 1.0,
				Conditions:  Conditions{}, // Match all
				Responses: []Response{
					{
						Type: ResponseTypeSuccess,
						Content: &SuccessResponse{
							MessageTemplate: bifrost.Ptr("Hello from {{provider}} using model {{model}}"),
						},
					},
				},
			},
		},
	}
	plugin, err := Init(config)
	if err != nil {
		t.Fatalf("Expected no error creating plugin, got: %v", err)
	}

	account := BaseAccount{}
	client, err := bifrost.Init(ctx, schemas.BifrostConfig{
		Account:    &account,
		LLMPlugins: []schemas.LLMPlugin{plugin},
		Logger:     bifrost.NewDefaultLogger(schemas.LogLevelError),
	})
	if err != nil {
		t.Fatalf("Error initializing Bifrost: %v", err)
	}
	defer client.Shutdown()

	response, bifrostErr := client.ChatCompletionRequest(ctx, &schemas.BifrostChatRequest{
		Provider: schemas.Anthropic,
		Model:    "claude-3",
		Input: []schemas.ChatMessage{
			{
				Role: schemas.ChatMessageRoleUser,
				Content: &schemas.ChatMessageContent{
					ContentStr: bifrost.Ptr("Hello, test message"),
				},
			},
		},
	})

	if bifrostErr != nil {
		t.Fatalf("Expected no error, got: %v", bifrostErr)
	}
	if response == nil {
		t.Fatal("Expected response")
	}
	if len(response.Choices) == 0 {
		t.Fatal("Expected at least one choice")
	}
	if response.Choices[0].ChatNonStreamResponseChoice.Message.Content.ContentStr == nil {
		t.Fatal("Expected content string")
	}
	expectedMessage := "Hello from anthropic using model claude-3"
	if *response.Choices[0].ChatNonStreamResponseChoice.Message.Content.ContentStr != expectedMessage {
		t.Errorf("Expected '%s', got: %s", expectedMessage, *response.Choices[0].ChatNonStreamResponseChoice.Message.Content.ContentStr)
	}
}

// TestMockerPlugin_Statistics tests plugin statistics tracking
func TestMockerPlugin_Statistics(t *testing.T) {
	ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
	config := MockerConfig{
		Enabled: true,
		Rules: []MockRule{
			{
				Name:        "stats-test",
				Enabled:     true,
				Priority:    100,
				Probability: 1.0,
				Conditions:  Conditions{}, // Match all
				Responses: []Response{
					{
						Type: ResponseTypeSuccess,
						Content: &SuccessResponse{
							Message: "Stats test response",
						},
					},
				},
			},
		},
	}
	plugin, err := Init(config)
	if err != nil {
		t.Fatalf("Expected no error creating plugin, got: %v", err)
	}

	account := BaseAccount{}
	client, err := bifrost.Init(ctx, schemas.BifrostConfig{
		Account:    &account,
		LLMPlugins: []schemas.LLMPlugin{plugin},
		Logger:     bifrost.NewDefaultLogger(schemas.LogLevelError),
	})
	if err != nil {
		t.Fatalf("Error initializing Bifrost: %v", err)
	}
	defer client.Shutdown()

	// Make multiple requests
	for i := 0; i < 3; i++ {
		_, _ = client.ChatCompletionRequest(ctx, &schemas.BifrostChatRequest{
			Provider: schemas.OpenAI,
			Model:    "gpt-4",
			Input: []schemas.ChatMessage{
				{
					Role: schemas.ChatMessageRoleUser,
					Content: &schemas.ChatMessageContent{
						ContentStr: bifrost.Ptr("Hello, test message"),
					},
				},
			},
		})
	}

	// Check statistics
	stats := plugin.GetStats()
	if stats.TotalRequests != 3 {
		t.Errorf("Expected 3 total requests, got %d", stats.TotalRequests)
	}
	if stats.MockedRequests != 3 {
		t.Errorf("Expected 3 mocked requests, got %d", stats.MockedRequests)
	}
	if stats.ResponsesGenerated != 3 {
		t.Errorf("Expected 3 responses generated, got %d", stats.ResponsesGenerated)
	}
	if stats.RuleHits["stats-test"] != 3 {
		t.Errorf("Expected 3 hits for 'stats-test' rule, got %d", stats.RuleHits["stats-test"])
	}
}

// TestMockerPlugin_ValidationErrors tests configuration validation
func TestMockerPlugin_ValidationErrors(t *testing.T) {
	tests := []struct {
		name        string
		config      MockerConfig
		expectError bool
	}{
		{
			name: "invalid default behavior",
			config: MockerConfig{
				Enabled:         true,
				DefaultBehavior: "invalid",
			},
			expectError: true,
		},
		{
			name: "missing rule name",
			config: MockerConfig{
				Enabled: true,
				Rules: []MockRule{
					{
						Name:    "", // Missing name
						Enabled: true,
						Responses: []Response{
							{
								Type: ResponseTypeSuccess,
								Content: &SuccessResponse{
									Message: "test",
								},
							},
						},
					},
				},
			},
			expectError: true,
		},
		{
			name: "invalid probability",
			config: MockerConfig{
				Enabled: true,
				Rules: []MockRule{
					{
						Name:        "test",
						Enabled:     true,
						Probability: 1.5, // Invalid probability > 1
						Responses: []Response{
							{
								Type: ResponseTypeSuccess,
								Content: &SuccessResponse{
									Message: "test",
								},
							},
						},
					},
				},
			},
			expectError: true,
		},
		{
			name: "valid configuration",
			config: MockerConfig{
				Enabled:         true,
				DefaultBehavior: DefaultBehaviorPassthrough,
				Rules: []MockRule{
					{
						Name:        "valid-rule",
						Enabled:     true,
						Probability: 0.5,
						Responses: []Response{
							{
								Type: ResponseTypeSuccess,
								Content: &SuccessResponse{
									Message: "Valid response",
								},
							},
						},
					},
				},
			},
			expectError: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			_, err := Init(tt.config)
			if tt.expectError && err == nil {
				t.Error("Expected error but got none")
			}
			if !tt.expectError && err != nil {
				t.Errorf("Expected no error but got: %v", err)
			}
		})
	}
}