bifrost/plugins/maxim/plugin_test.go

// Package maxim provides integration for Maxim's SDK as a Bifrost plugin.
// It includes tests for plugin initialization, Bifrost integration, and request/response tracing.
package maxim

import (
	"context"
	"fmt"
	"log"
	"os"
	"testing"

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

// getPlugin initializes and returns a Plugin instance for testing purposes.
// It sets up the Maxim logger with configuration from environment variables.
//
// Environment Variables:
//   - MAXIM_API_KEY: API key for Maxim SDK authentication
//   - MAXIM_LOG_REPO_ID: ID for the Maxim logger instance
//
// Returns:
//   - schemas.LLMPlugin: A configured plugin instance for request/response tracing
//   - error: Any error that occurred during plugin initialization
func getPlugin() (schemas.LLMPlugin, error) {
	// check if Maxim Logger variables are set
	if os.Getenv("MAXIM_API_KEY") == "" {
		return nil, fmt.Errorf("MAXIM_API_KEY is not set, please set it in your environment variables")
	}

	logger := bifrost.NewDefaultLogger(schemas.LogLevelDebug)
	plugin, err := Init(&Config{
		APIKey:    os.Getenv("MAXIM_API_KEY"),
		LogRepoID: os.Getenv("MAXIM_LOG_REPO_ID"),
	}, logger)
	if err != nil {
		return nil, err
	}

	return plugin, nil
}

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

// GetConfiguredProviders returns a list of supported providers for testing.
// Currently only supports OpenAI for simplicity in testing. You are free to add more providers as needed.
func (baseAccount *BaseAccount) GetConfiguredProviders() ([]schemas.ModelProvider, error) {
	return []schemas.ModelProvider{schemas.OpenAI}, nil
}

// GetKeysForProvider returns a mock API key configuration for testing.
// Uses the OPENAI_API_KEY environment variable for authentication.
func (baseAccount *BaseAccount) GetKeysForProvider(ctx context.Context, providerKey schemas.ModelProvider) ([]schemas.Key, error) {
	return []schemas.Key{
		{
			Value:  *schemas.NewEnvVar("env.OPENAI_API_KEY"),
			Models: []string{"gpt-4o-mini", "gpt-4-turbo"},
			Weight: 1.0,
		},
	}, nil
}

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

// TestMaximLoggerPlugin tests the integration of the Maxim Logger plugin with Bifrost.
// It performs the following steps:
// 1. Initializes the Maxim plugin with environment variables
// 2. Sets up a test Bifrost instance with the plugin
// 3. Makes a test chat completion request
//
// Required environment variables:
//   - MAXIM_API_KEY: Your Maxim API key
//   - MAXIM_LOGGER_ID: Your Maxim logger repository ID
//   - OPENAI_API_KEY: Your OpenAI API key for the test request
func TestMaximLoggerPlugin(t *testing.T) {
	ctx := context.Background()
	// Initialize the Maxim plugin
	plugin, err := getPlugin()
	if err != nil {
		t.Fatalf("Error setting up the plugin: %v", err)
	}

	account := BaseAccount{}

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

	// Make a test chat completion request
	_, bifrostErr := client.ChatCompletionRequest(schemas.NewBifrostContext(context.Background(), schemas.NoDeadline), &schemas.BifrostChatRequest{
		Provider: schemas.OpenAI,
		Model:    "gpt-4o-mini",
		Input: []schemas.ChatMessage{
			{
				Role: "user",
				Content: &schemas.ChatMessageContent{
					ContentStr: bifrost.Ptr("Hello, how are you?"),
				},
			},
		},
	})

	if bifrostErr != nil {
		log.Printf("Error in Bifrost request: %v", bifrostErr)
	}

	log.Println("Bifrost request completed, check your Maxim Dashboard for the trace")

	client.Shutdown()
}

// TestLogRepoIDSelection tests the single repository selection logic
func TestLogRepoIDSelection(t *testing.T) {
	tests := []struct {
		name         string
		defaultRepo  string
		headerRepo   string
		expectedRepo string
		shouldLog    bool
	}{
		{
			name:         "Header repo takes priority",
			defaultRepo:  "default-repo",
			headerRepo:   "header-repo",
			expectedRepo: "header-repo",
			shouldLog:    true,
		},
		{
			name:         "Fall back to default repo when no header",
			defaultRepo:  "default-repo",
			headerRepo:   "",
			expectedRepo: "default-repo",
			shouldLog:    true,
		},
		{
			name:         "Use header repo when no default",
			defaultRepo:  "",
			headerRepo:   "header-repo",
			expectedRepo: "header-repo",
			shouldLog:    true,
		},
		{
			name:         "Skip logging when neither available",
			defaultRepo:  "",
			headerRepo:   "",
			expectedRepo: "",
			shouldLog:    false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Create plugin with default repo
			plugin := &Plugin{
				defaultLogRepoID: tt.defaultRepo,
			}

			// Create context with header repo if provided
			ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
			if tt.headerRepo != "" {
				ctx.SetValue(LogRepoIDKey, tt.headerRepo)
			}

			// Test the selection logic
			result := plugin.getEffectiveLogRepoID(ctx)

			if result != tt.expectedRepo {
				t.Errorf("Expected repo '%s', got '%s'", tt.expectedRepo, result)
			}

			shouldLog := result != ""
			if shouldLog != tt.shouldLog {
				t.Errorf("Expected shouldLog=%t, got shouldLog=%t", tt.shouldLog, shouldLog)
			}
		})
	}
}

// TestPluginInitialization tests plugin initialization with different configs
func TestPluginInitialization(t *testing.T) {
	logger := bifrost.NewDefaultLogger(schemas.LogLevelDebug)
	tests := []struct {
		name        string
		config      Config
		expectError bool
	}{
		{
			name: "Valid config with both fields",
			config: Config{
				APIKey:    "test-api-key",
				LogRepoID: "test-repo-id",
			},
			expectError: false,
		},
		{
			name: "Valid config with only API key",
			config: Config{
				APIKey:    "test-api-key",
				LogRepoID: "",
			},
			expectError: false,
		},
		{
			name: "Invalid config - missing API key",
			config: Config{
				APIKey:    "",
				LogRepoID: "test-repo-id",
			},
			expectError: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Skip actual Maxim SDK initialization in tests
			if tt.expectError {
				_, err := Init(&tt.config, logger)
				if err == nil {
					t.Error("Expected error but got none")
				}
			} else {
				// For valid configs, we can't test actual initialization without real API key
				// Just test the validation logic
				if tt.config.APIKey == "" {
					t.Skip("Skipping valid config test - would need real Maxim API key")
				}
			}
		})
	}
}

// TestPluginName tests the plugin name functionality
func TestPluginName(t *testing.T) {
	plugin := &Plugin{}
	if plugin.GetName() != PluginName {
		t.Errorf("Expected plugin name '%s', got '%s'", PluginName, plugin.GetName())
	}
	if PluginName != "maxim" {
		t.Errorf("Expected PluginName constant to be 'maxim', got '%s'", PluginName)
	}
}