xinmi/internal/application/service/wiki_page_write_guard_test.go

package service

import (
	"context"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/application/repository"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/stretchr/testify/require"
)

type interveningWikiPageRepository struct {
	interfaces.WikiPageRepository
	beforeWrite func()
}

func (r *interveningWikiPageRepository) UpdateWithRevision(ctx context.Context, page *types.WikiPage, revision *types.WikiPageRevision) error {
	r.beforeWrite()
	return r.WikiPageRepository.UpdateWithRevision(ctx, page, revision)
}

func (r *interveningWikiPageRepository) UpdateMeta(ctx context.Context, page *types.WikiPage) error {
	r.beforeWrite()
	return r.WikiPageRepository.UpdateMeta(ctx, page)
}

func (r *interveningWikiPageRepository) UpdateAutoLinkedContent(ctx context.Context, page *types.WikiPage) error {
	r.beforeWrite()
	return r.WikiPageRepository.UpdateAutoLinkedContent(ctx, page)
}

func TestWikiGeneratedPageRejectsEditBeforeServiceRefetch(t *testing.T) {
	for _, metadataOnly := range []bool{false, true} {
		name := "content"
		if metadataOnly {
			name = "metadata"
		}
		t.Run(name, func(t *testing.T) {
			ctx, h, db := newWikiRevisionTestService(t)
			page, err := h.svc.CreatePage(ctx, &types.WikiPage{
				KnowledgeBaseID: "kb-guard", TenantID: 1, Slug: "concept/guard",
				Title: "Guard", PageType: types.WikiPageTypeConcept, Content: "original",
			})
			require.NoError(t, err)
			writeCtx := types.WithWikiPageWriteGuard(ctx, page)
			generated := *page
			generated.Content = "AI result based on old row"
			edit := *page
			if metadataOnly {
				edit.SortOrder = 7
			} else {
				edit.Content = "user revision"
			}
			latest, err := h.svc.UpdatePage(ctx, &edit)
			require.NoError(t, err)
			_, err = h.svc.UpdatePage(writeCtx, &generated)
			require.ErrorIs(t, err, repository.ErrWikiPageConflict)
			stored, err := h.svc.GetPageBySlug(ctx, page.KnowledgeBaseID, page.Slug)
			require.NoError(t, err)
			require.Equal(t, latest.Content, stored.Content)
			require.Equal(t, latest.SortOrder, stored.SortOrder)
			var count int64
			require.NoError(t, db.Model(&types.WikiPageRevision{}).Count(&count).Error)
			if metadataOnly {
				require.Zero(t, count)
			} else {
				require.EqualValues(t, 1, count)
			}
		})
	}
}

func TestWikiGeneratedPageRejectsEditBetweenRefetchAndWrite(t *testing.T) {
	for _, contentChanged := range []bool{false, true} {
		name := "metadata_write"
		if contentChanged {
			name = "content_write"
		}
		t.Run(name, func(t *testing.T) {
			ctx, h, db := newWikiRevisionTestService(t)
			page, err := h.svc.CreatePage(ctx, &types.WikiPage{
				KnowledgeBaseID: "kb-guard", TenantID: 1, Slug: "concept/guard",
				Title: "Guard", PageType: types.WikiPageTypeConcept, Content: "original",
			})
			require.NoError(t, err)
			repo := &interveningWikiPageRepository{
				WikiPageRepository: repository.NewWikiPageRepository(db),
				beforeWrite: func() {
					// Metadata changes keep the content version. They must still
					// invalidate the row used by an in-flight AI replacement.
					require.NoError(t, db.Model(&types.WikiPage{}).Where("id = ?", page.ID).
						Updates(map[string]any{"sort_order": 9, "updated_at": page.UpdatedAt.Add(time.Second)}).Error)
				},
			}
			svc := NewWikiPageService(repo, nil, nil, nil, nil)
			generated := *page
			generated.SourceRefs = types.StringArray{"new-source"}
			if contentChanged {
				generated.Content = "AI replacement"
			}
			_, err = svc.UpdatePage(types.WithWikiPageWriteGuard(ctx, page), &generated)
			require.ErrorIs(t, err, repository.ErrWikiPageConflict)
			stored, err := h.svc.GetPageBySlug(ctx, page.KnowledgeBaseID, page.Slug)
			require.NoError(t, err)
			require.Equal(t, "original", stored.Content)
			require.Equal(t, 9, stored.SortOrder)
			require.Equal(t, 1, stored.Version)
			require.Empty(t, stored.SourceRefs)
			var count int64
			require.NoError(t, db.Model(&types.WikiPageRevision{}).Count(&count).Error)
			require.Zero(t, count, "rejected writes must roll back their revision snapshot")
		})
	}
}

func TestWikiAutoLinksRejectConcurrentEdits(t *testing.T) {
	for _, scenario := range []struct {
		name          string
		beforeRefetch bool
		metadataOnly  bool
	}{
		{"content_before_refetch", true, false},
		{"metadata_before_refetch", true, true},
		{"content_before_write", false, false},
		{"metadata_before_write", false, true},
	} {
		t.Run(scenario.name, func(t *testing.T) {
			ctx, h, db := newWikiRevisionTestService(t)
			page, err := h.svc.CreatePage(ctx, &types.WikiPage{
				KnowledgeBaseID: "kb-guard", TenantID: 1, Slug: "concept/guard",
				Title: "Guard", PageType: types.WikiPageTypeConcept, Content: "original",
			})
			require.NoError(t, err)
			decorated := *page
			decorated.Content = "[[concept/other|original]]"
			var latest *types.WikiPage
			edit := func() {
				changed := *page
				changed.SourceRefs = types.StringArray{"new-source"}
				if !scenario.metadataOnly {
					changed.Content = "user revision"
				}
				latest, err = h.svc.UpdatePage(ctx, &changed)
				require.NoError(t, err)
			}
			svc := h.svc
			if scenario.beforeRefetch {
				edit()
			} else {
				svc = NewWikiPageService(&interveningWikiPageRepository{
					WikiPageRepository: repository.NewWikiPageRepository(db), beforeWrite: edit,
				}, nil, nil, nil, nil)
			}
			require.ErrorIs(t, svc.UpdateAutoLinkedContent(ctx, &decorated), repository.ErrWikiPageConflict)
			stored, err := h.svc.GetPageBySlug(ctx, page.KnowledgeBaseID, page.Slug)
			require.NoError(t, err)
			require.Equal(t, latest.Content, stored.Content)
			require.Equal(t, latest.SourceRefs, stored.SourceRefs)
			require.Equal(t, latest.Version, stored.Version)
			require.True(t, latest.UpdatedAt.Equal(stored.UpdatedAt))
			var revisions int64
			require.NoError(t, db.Model(&types.WikiPageRevision{}).Count(&revisions).Error)
			require.EqualValues(t, latest.Version-1, revisions)
		})
	}
}

func TestWikiMetadataRejectsStalePage(t *testing.T) {
	for _, archive := range []bool{false, true} {
		name := "new_source_refs"
		if archive {
			name = "archived"
		}
		t.Run(name, func(t *testing.T) {
			ctx, h, _ := newWikiRevisionTestService(t)
			page, err := h.svc.CreatePage(ctx, &types.WikiPage{
				KnowledgeBaseID: "kb-guard", TenantID: 1, Slug: "concept/guard",
				Title: "Guard", PageType: types.WikiPageTypeConcept, Content: "original",
			})
			require.NoError(t, err)
			changed := *page
			changed.SourceRefs = types.StringArray{"new-source"}
			changed.ChunkRefs = types.StringArray{"new-chunk"}
			if archive {
				changed.Status = types.WikiPageStatusArchived
			}
			latest, err := h.svc.UpdatePage(ctx, &changed)
			require.NoError(t, err)
			page.SortOrder = 7
			require.ErrorIs(t, h.svc.UpdatePageMeta(ctx, page), repository.ErrWikiPageConflict)
			stored, err := h.svc.GetPageBySlug(ctx, page.KnowledgeBaseID, page.Slug)
			require.NoError(t, err)
			require.Equal(t, latest.Status, stored.Status)
			require.Equal(t, latest.SourceRefs, stored.SourceRefs)
			require.Equal(t, latest.ChunkRefs, stored.ChunkRefs)
			require.Equal(t, latest.SortOrder, stored.SortOrder)
		})
	}
}

func TestWikiLinkMaintenancePreservesConcurrentMetadata(t *testing.T) {
	for _, remove := range []bool{false, true} {
		name := "add"
		if remove {
			name = "remove"
		}
		t.Run(name, func(t *testing.T) {
			ctx, h, db := newWikiRevisionTestService(t)
			page, err := h.svc.CreatePage(ctx, &types.WikiPage{
				KnowledgeBaseID: "kb-guard", TenantID: 1, Slug: "concept/guard",
				Title: "Guard", PageType: types.WikiPageTypeConcept, Content: "original",
				InLinks: types.StringArray{"concept/source"},
			})
			require.NoError(t, err)
			repo := &interveningWikiPageRepository{
				WikiPageRepository: repository.NewWikiPageRepository(db),
				beforeWrite: func() {
					require.NoError(t, db.Model(&types.WikiPage{}).Where("id = ?", page.ID).
						Updates(map[string]any{
							"source_refs": types.StringArray{"new-source"}, "chunk_refs": types.StringArray{"new-chunk"},
							"status": types.WikiPageStatusArchived, "updated_at": page.UpdatedAt.Add(time.Second),
						}).Error)
				},
			}
			svc := NewWikiPageService(repo, nil, nil, nil, nil).(*wikiPageService)
			if remove {
				svc.removeInLinks(ctx, page.KnowledgeBaseID, "concept/source", types.StringArray{page.Slug})
			} else {
				svc.updateInLinks(ctx, page.KnowledgeBaseID, "concept/another", types.StringArray{page.Slug})
			}
			stored, err := h.svc.GetPageBySlug(ctx, page.KnowledgeBaseID, page.Slug)
			require.NoError(t, err)
			require.Equal(t, types.WikiPageStatusArchived, stored.Status)
			require.Equal(t, types.StringArray{"new-source"}, stored.SourceRefs)
			require.Equal(t, types.StringArray{"new-chunk"}, stored.ChunkRefs)
			require.Equal(t, page.InLinks, stored.InLinks)
		})
	}
}

func TestWikiPublishPreservesConcurrentArchive(t *testing.T) {
	ctx, h, db := newWikiRevisionTestService(t)
	page, err := h.svc.CreatePage(ctx, &types.WikiPage{
		KnowledgeBaseID: "kb-guard", TenantID: 1, Slug: "concept/guard",
		Title: "Guard", PageType: types.WikiPageTypeConcept, Content: "original", Status: types.WikiPageStatusDraft,
	})
	require.NoError(t, err)
	repo := &interveningWikiPageRepository{
		WikiPageRepository: repository.NewWikiPageRepository(db),
		beforeWrite: func() {
			require.NoError(t, db.Model(&types.WikiPage{}).Where("id = ?", page.ID).
				Updates(map[string]any{"status": types.WikiPageStatusArchived, "updated_at": page.UpdatedAt.Add(time.Second)}).Error)
		},
	}
	svc := &wikiIngestService{wikiService: NewWikiPageService(repo, nil, nil, nil, nil)}
	failed := svc.publishDraftPages(ctx, page.KnowledgeBaseID, []string{page.Slug})
	require.ErrorIs(t, failed[page.Slug], repository.ErrWikiPageConflict)
	stored, err := h.svc.GetPageBySlug(ctx, page.KnowledgeBaseID, page.Slug)
	require.NoError(t, err)
	require.Equal(t, types.WikiPageStatusArchived, stored.Status)
}

func TestWikiAutoLinkWritesKeepSnapshotForCompensation(t *testing.T) {
	ctx, h, db := newWikiRevisionTestService(t)
	page, err := h.svc.CreatePage(ctx, &types.WikiPage{
		KnowledgeBaseID: "kb-guard", TenantID: 1, Slug: "concept/guard",
		Title: "Guard", PageType: types.WikiPageTypeConcept, Content: "original",
	})
	require.NoError(t, err)
	page.Content = "decorated"
	require.NoError(t, h.svc.UpdateAutoLinkedContent(ctx, page))
	page.Content = "original"
	require.NoError(t, h.svc.UpdateAutoLinkedContent(ctx, page))
	stored, err := h.svc.GetPageBySlug(ctx, page.KnowledgeBaseID, page.Slug)
	require.NoError(t, err)
	require.Equal(t, "original", stored.Content)
	require.Equal(t, 1, stored.Version)
	var revisions int64
	require.NoError(t, db.Model(&types.WikiPageRevision{}).Count(&revisions).Error)
	require.Zero(t, revisions)
}