first commit
This commit is contained in:
Beyhan Oğur
258
plugins/maxim/plugin_test.go
Normal file
258
plugins/maxim/plugin_test.go
Normal file
@@ -0,0 +1,258 @@
|
||||
// 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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user