gochat/backend/internal/repository/article_embedding_repo_test.go

package repository

import (
	"context"
	"testing"

	"github.com/pgvector/pgvector-go"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/gochat/gochat/internal/model"
)

func TestArticleEmbeddingRepo_DB(t *testing.T) {
	db := setupTestDB(t, &model.ArticleEmbedding{}, &model.Article{})
	repo := NewArticleEmbeddingRepo(db)
	assert.NotNil(t, repo.DB())
}

func TestArticleEmbeddingRepo_Upsert_Create(t *testing.T) {
	db := setupTestDB(t, &model.ArticleEmbedding{}, &model.Article{})
	repo := NewArticleEmbeddingRepo(db)

	account := &model.Account{Name: "AEOrg", Locale: "en", Active: true}
	require.NoError(t, db.Create(account).Error)
	portal := &model.Portal{AccountID: account.ID, Name: "Portal", Slug: "ae-portal"}
	require.NoError(t, db.Create(portal).Error)
	article := &model.Article{AccountID: account.ID, PortalID: portal.ID, Title: "Art", Slug: "ae-art", Status: "published"}
	require.NoError(t, db.Create(article).Error)

	emb := &model.ArticleEmbedding{
		ArticleID: article.ID,
		Term:      "test content",
	}
	err := repo.Upsert(context.Background(), emb)
	assert.NoError(t, err)
}

func TestArticleEmbeddingRepo_Upsert_Update(t *testing.T) {
	db := setupTestDB(t, &model.ArticleEmbedding{}, &model.Article{})
	repo := NewArticleEmbeddingRepo(db)

	account := &model.Account{Name: "AEUpdOrg", Locale: "en", Active: true}
	require.NoError(t, db.Create(account).Error)
	portal := &model.Portal{AccountID: account.ID, Name: "Portal", Slug: "ae-upd"}
	require.NoError(t, db.Create(portal).Error)
	article := &model.Article{AccountID: account.ID, PortalID: portal.ID, Title: "Art", Slug: "ae-upd-art", Status: "published"}
	require.NoError(t, db.Create(article).Error)

	// Create first
	emb := &model.ArticleEmbedding{ArticleID: article.ID, Term: "original"}
	require.NoError(t, repo.Upsert(context.Background(), emb))

	// Update
	emb2 := &model.ArticleEmbedding{ArticleID: article.ID, Term: "updated"}
	err := repo.Upsert(context.Background(), emb2)
	assert.NoError(t, err)

	found, _ := repo.GetByArticleID(context.Background(), article.ID)
	assert.Equal(t, "updated", found.Term)
}

func TestArticleEmbeddingRepo_GetByArticleID(t *testing.T) {
	db := setupTestDB(t, &model.ArticleEmbedding{}, &model.Article{})
	repo := NewArticleEmbeddingRepo(db)

	account := &model.Account{Name: "AEGetOrg", Locale: "en", Active: true}
	require.NoError(t, db.Create(account).Error)
	portal := &model.Portal{AccountID: account.ID, Name: "Portal", Slug: "ae-get"}
	require.NoError(t, db.Create(portal).Error)
	article := &model.Article{AccountID: account.ID, PortalID: portal.ID, Title: "Art", Slug: "ae-get-art", Status: "published"}
	require.NoError(t, db.Create(article).Error)

	emb := &model.ArticleEmbedding{ArticleID: article.ID, Term: "get test"}
	require.NoError(t, repo.Upsert(context.Background(), emb))

	found, err := repo.GetByArticleID(context.Background(), article.ID)
	assert.NoError(t, err)
	assert.Equal(t, "get test", found.Term)
}

func TestArticleEmbeddingRepo_GetByArticleID_NotFound(t *testing.T) {
	db := setupTestDB(t, &model.ArticleEmbedding{}, &model.Article{})
	repo := NewArticleEmbeddingRepo(db)

	_, err := repo.GetByArticleID(context.Background(), 9999)
	assert.Error(t, err)
}

func TestArticleEmbeddingRepo_DeleteByArticleID(t *testing.T) {
	db := setupTestDB(t, &model.ArticleEmbedding{}, &model.Article{})
	repo := NewArticleEmbeddingRepo(db)

	account := &model.Account{Name: "AEDelOrg", Locale: "en", Active: true}
	require.NoError(t, db.Create(account).Error)
	portal := &model.Portal{AccountID: account.ID, Name: "Portal", Slug: "ae-del"}
	require.NoError(t, db.Create(portal).Error)
	article := &model.Article{AccountID: account.ID, PortalID: portal.ID, Title: "Art", Slug: "ae-del-art", Status: "published"}
	require.NoError(t, db.Create(article).Error)

	emb := &model.ArticleEmbedding{ArticleID: article.ID, Term: "del test"}
	require.NoError(t, repo.Upsert(context.Background(), emb))

	err := repo.DeleteByArticleID(context.Background(), article.ID)
	assert.NoError(t, err)

	_, err = repo.GetByArticleID(context.Background(), article.ID)
	assert.Error(t, err)
}

func TestArticleEmbeddingRepo_SearchByEmbedding_ExcludesSoftDeletedArticleAndEmbedding(t *testing.T) {
	skipIfSQLite(t)

	db := setupTestDB(t, &model.ArticleEmbedding{}, &model.Article{}, &model.Portal{})
	repo := NewArticleEmbeddingRepo(db)
	account := &model.Account{Name: "Soft Delete Search", Locale: "en", Active: true}
	require.NoError(t, db.Create(account).Error)
	portal := &model.Portal{AccountID: account.ID, Name: "Portal", Slug: "soft-delete-search"}
	require.NoError(t, db.Create(portal).Error)

	embedding := pgvector.NewVector([]float32{1, 0})
	articles := make([]model.Article, 3)
	for i, name := range []string{"live", "deleted-article", "deleted-embedding"} {
		articles[i] = model.Article{
			AccountID: account.ID,
			PortalID:  portal.ID,
			Title:     name,
			Slug:      "soft-delete-search-" + name,
			Status:    "published",
		}
		require.NoError(t, db.Create(&articles[i]).Error)
		require.NoError(t, repo.Upsert(context.Background(), &model.ArticleEmbedding{
			ArticleID:       articles[i].ID,
			VectorEmbedding: embedding,
			Term:            name,
		}))
	}

	before, err := repo.SearchByEmbedding(context.Background(), portal.ID, embedding, 10)
	require.NoError(t, err)
	require.Len(t, before, 3)
	require.NotNil(t, before[0].SemanticDistance)
	require.NoError(t, db.Delete(&articles[1]).Error)
	require.NoError(t, repo.DeleteByArticleID(context.Background(), articles[2].ID))

	after, err := repo.SearchByEmbedding(context.Background(), portal.ID, embedding, 10)
	require.NoError(t, err)
	require.Len(t, after, 1)
	assert.Equal(t, articles[0].ID, after[0].ID)
}