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 }