vet/agent/memory_test.go

package agent

import (
	"context"
	"fmt"
	"sync"
	"testing"

	"github.com/cloudwego/eino/schema"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestNewSimpleMemory(t *testing.T) {
	memory, err := NewSimpleMemory()
	require.NoError(t, err)
	assert.NotNil(t, memory)

	// Test that initial interactions are empty
	ctx := context.Background()
	interactions, err := memory.GetInteractions(ctx)
	require.NoError(t, err)
	assert.Equal(t, 0, len(interactions))
}

func TestSimpleMemory_AddInteraction(t *testing.T) {
	memory, err := NewSimpleMemory()
	require.NoError(t, err)

	ctx := context.Background()
	message := &schema.Message{
		Role:    schema.User,
		Content: "test message",
	}

	err = memory.AddInteraction(ctx, message)
	assert.NoError(t, err)

	interactions, err := memory.GetInteractions(ctx)
	require.NoError(t, err)
	assert.Equal(t, 1, len(interactions))
	assert.Equal(t, message, interactions[0])
}

func TestSimpleMemory_AddMultipleInteractions(t *testing.T) {
	memory, err := NewSimpleMemory()
	require.NoError(t, err)

	ctx := context.Background()
	messages := []*schema.Message{
		{Role: schema.User, Content: "first message"},
		{Role: schema.Assistant, Content: "second message"},
		{Role: schema.User, Content: "third message"},
	}

	for _, msg := range messages {
		err = memory.AddInteraction(ctx, msg)
		assert.NoError(t, err)
	}

	interactions, err := memory.GetInteractions(ctx)
	require.NoError(t, err)
	assert.Equal(t, 3, len(interactions))

	for i, msg := range messages {
		assert.Equal(t, msg, interactions[i])
	}
}

func TestSimpleMemory_GetInteractions(t *testing.T) {
	memory, err := NewSimpleMemory()
	require.NoError(t, err)

	ctx := context.Background()

	// Test empty interactions
	interactions, err := memory.GetInteractions(ctx)
	assert.NoError(t, err)
	assert.NotNil(t, interactions)
	assert.Equal(t, 0, len(interactions))

	// Add some interactions
	message1 := &schema.Message{Role: schema.User, Content: "message 1"}
	message2 := &schema.Message{Role: schema.Assistant, Content: "message 2"}

	err = memory.AddInteraction(ctx, message1)
	require.NoError(t, err)
	err = memory.AddInteraction(ctx, message2)
	require.NoError(t, err)

	interactions, err = memory.GetInteractions(ctx)
	assert.NoError(t, err)
	assert.Equal(t, 2, len(interactions))
	assert.Equal(t, message1, interactions[0])
	assert.Equal(t, message2, interactions[1])
}

func TestSimpleMemory_Clear(t *testing.T) {
	memory, err := NewSimpleMemory()
	require.NoError(t, err)

	ctx := context.Background()

	// Add some interactions
	message1 := &schema.Message{Role: schema.User, Content: "message 1"}
	message2 := &schema.Message{Role: schema.Assistant, Content: "message 2"}

	err = memory.AddInteraction(ctx, message1)
	require.NoError(t, err)
	err = memory.AddInteraction(ctx, message2)
	require.NoError(t, err)

	// Verify interactions exist
	interactions, err := memory.GetInteractions(ctx)
	require.NoError(t, err)
	assert.Equal(t, 2, len(interactions))

	// Clear interactions
	err = memory.Clear(ctx)
	assert.NoError(t, err)

	// Verify interactions are cleared
	interactions, err = memory.GetInteractions(ctx)
	assert.NoError(t, err)
	assert.Equal(t, 0, len(interactions))
}

func TestSimpleMemory_NilInteraction(t *testing.T) {
	memory, err := NewSimpleMemory()
	require.NoError(t, err)

	ctx := context.Background()

	// Test adding nil interaction
	err = memory.AddInteraction(ctx, nil)
	assert.NoError(t, err)

	interactions, err := memory.GetInteractions(ctx)
	require.NoError(t, err)
	assert.Equal(t, 1, len(interactions))
	assert.Nil(t, interactions[0])
}

func TestSimpleMemory_ConcurrentAccess(t *testing.T) {
	memory, err := NewSimpleMemory()
	require.NoError(t, err)

	ctx := context.Background()
	numGoroutines := 100
	messagesPerGoroutine := 10

	var wg sync.WaitGroup
	wg.Add(numGoroutines)

	// Concurrent writes
	for i := 0; i < numGoroutines; i++ {
		go func(goroutineID int) {
			defer wg.Done()
			for j := 0; j < messagesPerGoroutine; j++ {
				message := &schema.Message{
					Role:    schema.User,
					Content: fmt.Sprintf("goroutine-%d-message-%d", goroutineID, j),
				}
				err := memory.AddInteraction(ctx, message)
				assert.NoError(t, err)
			}
		}(i)
	}

	wg.Wait()

	// Verify all interactions were added
	interactions, err := memory.GetInteractions(ctx)
	require.NoError(t, err)
	assert.Equal(t, numGoroutines*messagesPerGoroutine, len(interactions))
}

func TestSimpleMemory_ConcurrentReadWrite(t *testing.T) {
	memory, err := NewSimpleMemory()
	require.NoError(t, err)

	ctx := context.Background()
	numReaders := 10
	numWriters := 10
	messagesPerWriter := 5

	var wg sync.WaitGroup
	wg.Add(numReaders + numWriters)

	// Concurrent writers
	for i := 0; i < numWriters; i++ {
		go func(writerID int) {
			defer wg.Done()
			for j := 0; j < messagesPerWriter; j++ {
				message := &schema.Message{
					Role:    schema.User,
					Content: fmt.Sprintf("writer-%d-message-%d", writerID, j),
				}
				err := memory.AddInteraction(ctx, message)
				assert.NoError(t, err)
			}
		}(i)
	}

	// Concurrent readers
	for i := 0; i < numReaders; i++ {
		go func() {
			defer wg.Done()
			for j := 0; j < messagesPerWriter; j++ {
				interactions, err := memory.GetInteractions(ctx)
				assert.NoError(t, err)
				assert.NotNil(t, interactions)
				// Length can vary due to concurrent writes
				assert.GreaterOrEqual(t, len(interactions), 0)
			}
		}()
	}

	wg.Wait()

	// Final verification
	interactions, err := memory.GetInteractions(ctx)
	require.NoError(t, err)
	assert.Equal(t, numWriters*messagesPerWriter, len(interactions))
}

func TestSimpleMemory_ClearDuringConcurrentAccess(t *testing.T) {
	memory, err := NewSimpleMemory()
	require.NoError(t, err)

	ctx := context.Background()
	numWriters := 5
	messagesPerWriter := 10

	var wg sync.WaitGroup
	wg.Add(numWriters + 1) // +1 for the clearer

	// Add some initial interactions
	for i := 0; i < 5; i++ {
		message := &schema.Message{
			Role:    schema.User,
			Content: fmt.Sprintf("initial-message-%d", i),
		}
		err := memory.AddInteraction(ctx, message)
		require.NoError(t, err)
	}

	// Concurrent writers
	for i := 0; i < numWriters; i++ {
		go func(writerID int) {
			defer wg.Done()
			for j := 0; j < messagesPerWriter; j++ {
				message := &schema.Message{
					Role:    schema.User,
					Content: fmt.Sprintf("writer-%d-message-%d", writerID, j),
				}
				err := memory.AddInteraction(ctx, message)
				assert.NoError(t, err)
			}
		}(i)
	}

	// Clear operation
	go func() {
		defer wg.Done()
		// Clear after some writes have happened
		err := memory.Clear(ctx)
		assert.NoError(t, err)
	}()

	wg.Wait()

	// Final state check - should be consistent
	interactions, err := memory.GetInteractions(ctx)
	require.NoError(t, err)
	assert.NotNil(t, interactions)
	// The exact number depends on timing of clear operation
	assert.GreaterOrEqual(t, len(interactions), 0)
}