Files
agent/backend/internal/initdb/dataset_test.go

97 lines
3.0 KiB
Go

package initdb
import (
"fmt"
"testing"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
"senlinai-agent/backend/internal/models"
)
func TestInitDatasetCreatesDefaultSourcesWhenOwnerHasNoSources(t *testing.T) {
database := newDatasetInitTestDatabase(t)
root, err := InitUser(database)
require.NoError(t, err)
require.NoError(t, InitDataset(database, root.ID, root.Identity))
require.NoError(t, InitDataset(database, root.ID, root.Identity))
var sources []models.SaDatasetSource
require.NoError(t, database.Order("id asc").Find(&sources).Error)
require.Len(t, sources, len(defaultDatasetSources))
for index, expected := range defaultDatasetSources {
require.Equal(t, root.ID, sources[index].OwnerID)
require.Equal(t, root.Identity, sources[index].OwnerIdentity)
require.Equal(t, expected.Name, sources[index].Name)
require.Equal(t, expected.URL, sources[index].URL)
require.Equal(t, expected.IconURL, sources[index].IconURL)
require.NotNil(t, sources[index].SeedKey)
require.Equal(t, expected.Key, *sources[index].SeedKey)
require.Equal(t, "rss", sources[index].Kind)
require.True(t, sources[index].Enabled)
}
}
func TestInitDatasetDoesNothingWhenOwnerAlreadyHasSource(t *testing.T) {
database := newDatasetInitTestDatabase(t)
root, err := InitUser(database)
require.NoError(t, err)
existing := models.SaDatasetSource{
OwnerID: root.ID,
Name: "Existing source",
Kind: "manual",
Enabled: true,
}
require.NoError(t, database.Create(&existing).Error)
require.NoError(t, InitDataset(database, root.ID, root.Identity))
var sources []models.SaDatasetSource
require.NoError(t, database.Find(&sources).Error)
require.Len(t, sources, 1)
require.Equal(t, existing.Identity, sources[0].Identity)
}
func TestInitDatasetCountsSourcesByOwner(t *testing.T) {
database := newDatasetInitTestDatabase(t)
root, err := InitUser(database)
require.NoError(t, err)
other := models.SaUser{
Email: "other@example.com", DisplayName: "Other",
PasswordHash: "hash", Role: "user",
}
require.NoError(t, database.Create(&other).Error)
require.NoError(t, database.Create(&models.SaDatasetSource{
OwnerID: other.ID, Name: "Other source", Kind: "manual", Enabled: true,
}).Error)
require.NoError(t, InitDataset(database, root.ID, root.Identity))
var rootCount int64
require.NoError(t, database.Model(&models.SaDatasetSource{}).
Where("owner_id = ?", root.ID).
Count(&rootCount).Error)
require.Equal(t, int64(len(defaultDatasetSources)), rootCount)
var otherCount int64
require.NoError(t, database.Model(&models.SaDatasetSource{}).
Where("owner_id = ?", other.ID).
Count(&otherCount).Error)
require.Equal(t, int64(1), otherCount)
}
func newDatasetInitTestDatabase(t *testing.T) *gorm.DB {
t.Helper()
database, err := gorm.Open(
sqlite.Open(fmt.Sprintf("file:%s?mode=memory&cache=shared", t.Name())),
&gorm.Config{},
)
require.NoError(t, err)
require.NoError(t, database.AutoMigrate(
&models.SaUser{},
&models.SaDatasetSource{},
))
return database
}