agent/backend/internal/initdb/dataset_test.go

package initdb

import (
	"fmt"
	"testing"
	"time"

	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
	"senlinai-agent/backend/internal/models"
)

func TestInitDatasetCreatesDefaultSourcesWhenTableIsEmpty(t *testing.T) {
	database := newDatasetInitTestDatabase(t)
	require.NoError(t, InitUser(database))

	require.NoError(t, InitDataset(database))
	require.NoError(t, InitDataset(database))

	var root models.SaUser
	require.NoError(t, database.Where("email = ?", rootUsername).First(&root).Error)
	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, 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)
		require.Equal(t, root.ID, sources[index].OwnerID)
		require.Equal(t, root.Identity, sources[index].OwnerIdentity)
		require.Equal(t, root.ID, sources[index].CreatedBy)
		require.Equal(t, root.Identity, sources[index].CreatedByIdentity)
	}
}

func TestInitDatasetAddsMissingDefaultsWithoutChangingCustomSources(t *testing.T) {
	database := newDatasetInitTestDatabase(t)
	require.NoError(t, InitUser(database))
	var root models.SaUser
	require.NoError(t, database.Where("email = ?", rootUsername).First(&root).Error)
	existing := models.SaDatasetSource{
		CreatedBy: root.ID,
		Name:      "Existing source",
		Kind:      "manual",
		Enabled:   true,
	}
	require.NoError(t, database.Create(&existing).Error)

	require.NoError(t, InitDataset(database))

	var sources []models.SaDatasetSource
	require.NoError(t, database.Find(&sources).Error)
	require.Len(t, sources, len(defaultDatasetSources)+1)
	var preserved models.SaDatasetSource
	require.NoError(t, database.Where("identity = ?", existing.Identity).First(&preserved).Error)
	require.Equal(t, existing.Name, preserved.Name)
	require.Nil(t, preserved.SeedKey)
}

func TestInitDatasetBackfillsLegacySourceOwners(t *testing.T) {
	database := newDatasetInitTestDatabase(t)
	require.NoError(t, InitUser(database))
	other := models.SaUser{
		Email: "other@example.com", DisplayName: "Other",
		PasswordHash: "hash", Role: "user",
	}
	require.NoError(t, database.Create(&other).Error)
	source := models.SaDatasetSource{
		CreatedBy: other.ID, Name: "Existing source", Kind: "manual", Enabled: true,
	}
	require.NoError(t, database.Create(&source).Error)
	require.NoError(t, database.Model(&source).Updates(map[string]any{
		"owner_id": 0, "owner_identity": "",
	}).Error)

	require.NoError(t, InitDataset(database))

	require.NoError(t, database.First(&source, source.ID).Error)
	require.Equal(t, other.ID, source.OwnerID)
	require.Equal(t, other.Identity, source.OwnerIdentity)
	require.Equal(t, other.ID, source.CreatedBy)
}

func TestInitDatasetRejectsMissingRootInsteadOfAssigningDefaultsToAnotherUser(t *testing.T) {
	database := newDatasetInitTestDatabase(t)
	user := models.SaUser{
		Email:        "existing@example.com",
		DisplayName:  "Existing",
		PasswordHash: "hash",
		Role:         "user",
	}
	require.NoError(t, database.Create(&user).Error)

	require.Error(t, InitDataset(database))

	var sources []models.SaDatasetSource
	require.NoError(t, database.Find(&sources).Error)
	require.Empty(t, sources)
}

func TestInitDatasetRepairsLegacyRSSHubURLsForRootOnly(t *testing.T) {
	database := newDatasetInitTestDatabase(t)
	require.NoError(t, InitUser(database))
	var root models.SaUser
	require.NoError(t, database.Where("email = ?", rootUsername).First(&root).Error)
	other := models.SaUser{
		Email: "other@example.com", DisplayName: "Other",
		PasswordHash: "hash", Role: "user",
	}
	require.NoError(t, database.Create(&other).Error)

	rootLegacy := models.SaDatasetSource{
		CreatedBy: root.ID, Name: "财联社-热门", Kind: "rss",
		URL: "https://rsshub.app/cls/hot", Enabled: true,
	}
	otherLegacy := models.SaDatasetSource{
		CreatedBy: other.ID, Name: "财联社-热门", Kind: "rss",
		URL: "https://rsshub.app/cls/hot", Enabled: true,
	}
	require.NoError(t, database.Create(&rootLegacy).Error)
	require.NoError(t, database.Create(&otherLegacy).Error)

	require.NoError(t, InitDataset(database))

	require.NoError(t, database.First(&rootLegacy, rootLegacy.ID).Error)
	require.Equal(t, "https://rsshub.ktachibana.party/cls/hot", rootLegacy.URL)
	require.NotNil(t, rootLegacy.SeedKey)
	require.Equal(t, "rsshub-cls-hot", *rootLegacy.SeedKey)
	require.NoError(t, database.First(&otherLegacy, otherLegacy.ID).Error)
	require.Equal(t, "https://rsshub.app/cls/hot", otherLegacy.URL)
	require.Nil(t, otherLegacy.SeedKey)
}

func TestInitDatasetPreparesCollectedItemUniqueness(t *testing.T) {
	database := newDatasetInitTestDatabase(t)
	require.NoError(t, InitUser(database))
	var root models.SaUser
	require.NoError(t, database.Where("email = ?", rootUsername).First(&root).Error)
	source := models.SaDatasetSource{
		CreatedBy: root.ID, Name: "Existing feed", Kind: "rss",
		URL: "https://example.com/feed", Enabled: true,
	}
	require.NoError(t, database.Create(&source).Error)
	now := time.Now()
	for index, externalID := range []string{"duplicate", "duplicate", "", ""} {
		require.NoError(t, database.Exec(`
			INSERT INTO sa_dataset_items
				(identity, source_id, created_by, external_id, title, status, created_at, updated_at)
			VALUES (?, ?, ?, ?, ?, ?, ?, ?)
		`, fmt.Sprintf("item-%d", index), source.ID, root.ID, externalID,
			fmt.Sprintf("Item %d", index), "unread", now, now).Error)
	}

	require.NoError(t, InitDataset(database))

	var collectedCount int64
	require.NoError(t, database.Model(&models.SaDatasetItem{}).
		Where("source_id = ? AND external_id = ?", source.ID, "duplicate").
		Count(&collectedCount).Error)
	require.Equal(t, int64(1), collectedCount)
	var manualCount int64
	require.NoError(t, database.Model(&models.SaDatasetItem{}).
		Where("source_id = ? AND external_id IS NULL", source.ID).
		Count(&manualCount).Error)
	require.Equal(t, int64(2), manualCount)
	require.Error(t, database.Exec(`
		INSERT INTO sa_dataset_items
			(identity, source_id, created_by, external_id, title, status, created_at, updated_at)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?)
	`, "item-new", source.ID, root.ID, "duplicate", "Duplicate", "unread", now, now).Error)
}

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{},
		&models.SaDatasetItem{},
	))
	return database
}