summaryrefslogtreecommitdiff
path: root/storage/storage_test.go
diff options
context:
space:
mode:
authorGrail Finder <wohilas@gmail.com>2024-11-20 20:47:49 +0300
committerGrail Finder <wohilas@gmail.com>2024-11-20 20:47:49 +0300
commit5ccad20bd680dc443b30f0decc8fca13427dc70d (patch)
tree5ed20ce680c09609a29f880cf41c976301f3e031 /storage/storage_test.go
parentfc517c2c69d96501f1adc5a021b39b9eff22e4d7 (diff)
Feat: add memory [wip]
Diffstat (limited to 'storage/storage_test.go')
-rw-r--r--storage/storage_test.go77
1 files changed, 75 insertions, 2 deletions
diff --git a/storage/storage_test.go b/storage/storage_test.go
index 0bf1fd6..ad1f1bf 100644
--- a/storage/storage_test.go
+++ b/storage/storage_test.go
@@ -2,6 +2,9 @@ package storage
import (
"elefant/models"
+ "fmt"
+ "log/slog"
+ "os"
"testing"
"time"
@@ -9,6 +12,76 @@ import (
"github.com/jmoiron/sqlx"
)
+func TestMemories(t *testing.T) {
+ db, err := sqlx.Open("sqlite", ":memory:")
+ if err != nil {
+ t.Fatalf("Failed to open SQLite in-memory database: %v", err)
+ }
+ defer db.Close()
+ _, err = db.Exec(`
+CREATE TABLE IF NOT EXISTS memories (
+ agent TEXT NOT NULL,
+ topic TEXT NOT NULL,
+ mind TEXT NOT NULL,
+ created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ PRIMARY KEY (agent, topic)
+);`)
+ if err != nil {
+ t.Fatalf("Failed to create chat table: %v", err)
+ }
+ provider := ProviderSQL{
+ db: db,
+ logger: slog.New(slog.NewJSONHandler(os.Stdout, nil)),
+ }
+ // Create a sample memory for testing
+ sampleMemory := &models.Memory{
+ Agent: "testAgent",
+ Topic: "testTopic",
+ Mind: "testMind",
+ CreatedAt: time.Now(),
+ UpdatedAt: time.Now(),
+ }
+ cases := []struct {
+ memory *models.Memory
+ }{
+ {memory: sampleMemory},
+ }
+ for i, tc := range cases {
+ t.Run(fmt.Sprintf("run_%d", i), func(t *testing.T) {
+ // Recall topics: get no rows
+ topics, err := provider.RecallTopics(tc.memory.Agent)
+ if err != nil {
+ t.Fatalf("Failed to recall topics: %v", err)
+ }
+ if len(topics) != 0 {
+ t.Fatalf("Expected no topics, got: %v", topics)
+ }
+ // Memorise
+ _, err = provider.Memorise(tc.memory)
+ if err != nil {
+ t.Fatalf("Failed to memorise: %v", err)
+ }
+ // Recall topics: has topics
+ topics, err = provider.RecallTopics(tc.memory.Agent)
+ if err != nil {
+ t.Fatalf("Failed to recall topics: %v", err)
+ }
+ if len(topics) == 0 {
+ t.Fatalf("Expected topics, got none")
+ }
+ // Recall
+ content, err := provider.Recall(tc.memory.Agent, tc.memory.Topic)
+ if err != nil {
+ t.Fatalf("Failed to recall: %v", err)
+ }
+ if content != tc.memory.Mind {
+ t.Fatalf("Expected content: %v, got: %v", tc.memory.Mind, content)
+ }
+ })
+ }
+}
+
func TestChatHistory(t *testing.T) {
// Create an in-memory SQLite database
db, err := sqlx.Open("sqlite", ":memory:")
@@ -18,10 +91,10 @@ func TestChatHistory(t *testing.T) {
defer db.Close()
// Create the chat table
_, err = db.Exec(`
- CREATE TABLE chat (
+ CREATE TABLE chats (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT NOT NULL,
- msgs TEXT NOT NULL, -- Store messages as a comma-separated string
+ msgs TEXT NOT NULL,
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
);`)