mirror of
https://github.com/mariocandela/beelzebub.git
synced 2025-07-01 18:47:26 -04:00
d677cd20b9
* Feat: Add support for logging which "command" was matched for SSH and HTTP strategies. * Feat: Convert to precompiling regexp at config load time. This allows for errors to be presented to the user during startup, and provides better performance for complex regexp. * Feat:Bump Golang version to latest stable 1.24 * Feat: Add a cleanup routine for HistoryStore, default TTL for events is 1 hour since last interaction. * Feat: Add new command line flag "memLimitMiB" with a default value of 100. --------- Signed-off-by: Bryan Nolen <bryan@arc.net.au> Signed-off-by: Mario Candela <mario.candela.personal@gmail.com> Co-authored-by: Mario Candela <mario.candela.personal@gmail.com>
67 lines
2.0 KiB
Go
67 lines
2.0 KiB
Go
package historystore
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/mariocandela/beelzebub/v3/plugins"
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
func TestNewHistoryStore(t *testing.T) {
|
|
hs := NewHistoryStore()
|
|
assert.NotNil(t, hs)
|
|
assert.NotNil(t, hs.sessions)
|
|
}
|
|
|
|
func TestHasKey(t *testing.T) {
|
|
hs := NewHistoryStore()
|
|
hs.sessions["testKey"] = HistoryEvent{Messages: []plugins.Message{}}
|
|
assert.True(t, hs.HasKey("testKey"))
|
|
assert.False(t, hs.HasKey("nonExistentKey"))
|
|
}
|
|
|
|
func TestQuery(t *testing.T) {
|
|
hs := NewHistoryStore()
|
|
expectedMessages := []plugins.Message{{Role: "user", Content: "Hello"}}
|
|
hs.sessions["testKey"] = HistoryEvent{Messages: expectedMessages}
|
|
actualMessages := hs.Query("testKey")
|
|
assert.Equal(t, expectedMessages, actualMessages)
|
|
}
|
|
|
|
func TestAppend(t *testing.T) {
|
|
hs := NewHistoryStore()
|
|
message1 := plugins.Message{Role: "user", Content: "Hello"}
|
|
message2 := plugins.Message{Role: "assistant", Content: "Hi"}
|
|
hs.Append("testKey", message1)
|
|
assert.Equal(t, []plugins.Message{message1}, hs.sessions["testKey"].Messages)
|
|
hs.Append("testKey", message2)
|
|
assert.Equal(t, []plugins.Message{message1, message2}, hs.sessions["testKey"].Messages)
|
|
}
|
|
|
|
func TestAppendNilSessions(t *testing.T) {
|
|
hs := &HistoryStore{}
|
|
message1 := plugins.Message{Role: "user", Content: "Hello"}
|
|
hs.Append("testKey", message1)
|
|
assert.NotNil(t, hs.sessions)
|
|
assert.Equal(t, []plugins.Message{message1}, hs.sessions["testKey"].Messages)
|
|
}
|
|
|
|
func TestHistoryCleaner(t *testing.T) {
|
|
hs := NewHistoryStore()
|
|
hs.Append("testKey", plugins.Message{Role: "user", Content: "Hello"})
|
|
hs.Append("testKey2", plugins.Message{Role: "user", Content: "Hello"})
|
|
|
|
// Make key older than MaxHistoryAge
|
|
e := hs.sessions["testKey"]
|
|
e.LastSeen = time.Now().Add(-MaxHistoryAge * 2)
|
|
hs.sessions["testKey"] = e
|
|
|
|
CleanerInterval = 5 * time.Second // Override for the test.
|
|
hs.HistoryCleaner()
|
|
time.Sleep(CleanerInterval + (1 * time.Second))
|
|
|
|
assert.False(t, hs.HasKey("testKey"))
|
|
assert.True(t, hs.HasKey("testKey2"))
|
|
}
|