FranklinWH-Toolbox/service/utils/agentserver/guard_test.go

package agentserver

import (
	"context"
	"errors"
	"testing"
)

type fakeStore struct {
	owners    map[string]string // sessionId → userId（未删除）
	resources map[string]string // resourceId → sessionId
	err       error
}

func (f *fakeStore) SessionOwner(_ context.Context, sid string) (string, bool, error) {
	if f.err != nil {
		return "", false, f.err
	}
	o, ok := f.owners[sid]
	return o, ok, nil
}

func (f *fakeStore) ResourceSession(_ context.Context, id string) (string, bool, error) {
	s, ok := f.resources[id]
	return s, ok, nil
}

func (f *fakeStore) RecentSessionIds(_ context.Context, uid string, _ int) ([]string, error) {
	var out []string
	for s, o := range f.owners {
		if o == uid {
			out = append(out, s)
		}
	}
	return out, nil
}

type fakeUpstream struct {
	runs  map[string]string // runId → sessionId
	files map[string][]string
	itx   map[string][]string
}

func (f *fakeUpstream) GetRun(_ context.Context, id string) (*RunDto, error) {
	s, ok := f.runs[id]
	if !ok {
		return nil, ErrNotFound
	}
	return &RunDto{Id: id, SessionId: s}, nil
}

func (f *fakeUpstream) ListFiles(_ context.Context, sid string) ([]FileDto, error) {
	var out []FileDto
	for _, id := range f.files[sid] {
		out = append(out, FileDto{Id: id, SessionId: sid})
	}
	return out, nil
}

func (f *fakeUpstream) PendingInteractions(_ context.Context, sid string) ([]InteractionRefDto, error) {
	var out []InteractionRefDto
	for _, id := range f.itx[sid] {
		out = append(out, InteractionRefDto{InteractionId: id})
	}
	return out, nil
}

type fakeSink struct{ got []Resource }

func (f *fakeSink) Record(rs []Resource) { f.got = append(f.got, rs...) }

func newTestGuard() (*Guard, *fakeSink) {
	sink := &fakeSink{}
	return &Guard{
		Store: &fakeStore{
			owners:    map[string]string{"ses_a1": "alice", "ses_b1": "bob"},
			resources: map[string]string{"run_a": "ses_a1", "file_b": "ses_b1"},
		},
		Upstream: &fakeUpstream{
			runs:  map[string]string{"run_a2": "ses_a1", "run_b2": "ses_b1"},
			files: map[string][]string{"ses_a1": {"file_a2"}, "ses_b1": {"file_b2"}},
			itx:   map[string][]string{"ses_a1": {"itx_a2"}, "ses_b1": {"itx_b2"}},
		},
		Sink: sink,
	}, sink
}

func TestGuardCheck(t *testing.T) {
	g, sink := newTestGuard()
	ctx := context.Background()
	cases := []struct {
		name  string
		route Route
		ok    bool
		sid   string
	}{
		{"元数据", Route{Owner: OwnerNone}, true, ""},
		{"本人会话", Route{Owner: OwnerSession, Id: "ses_a1"}, true, "ses_a1"},
		{"他人会话", Route{Owner: OwnerSession, Id: "ses_b1"}, false, ""},
		{"不存在的会话", Route{Owner: OwnerSession, Id: "ses_zz"}, false, ""},
		{"查询参数会话（他人）", Route{Owner: OwnerQuerySession, Id: "ses_b1"}, false, ""},
		{"索引中的本人 run", Route{Owner: OwnerRun, Id: "run_a"}, true, "ses_a1"},
		{"索引中的他人文件", Route{Owner: OwnerFile, Id: "file_b"}, false, ""},
		{"反查本人 run", Route{Owner: OwnerRun, Id: "run_a2"}, true, "ses_a1"},
		{"反查他人 run", Route{Owner: OwnerRun, Id: "run_b2"}, false, ""},
		{"反查不存在的 run", Route{Owner: OwnerRun, Id: "run_zz"}, false, ""},
		{"反查本人文件", Route{Owner: OwnerFile, Id: "file_a2"}, true, "ses_a1"},
		{"反查他人文件", Route{Owner: OwnerFile, Id: "file_b2"}, false, ""},
		{"反查本人交互", Route{Owner: OwnerInteraction, Id: "itx_a2"}, true, "ses_a1"},
		{"反查他人交互", Route{Owner: OwnerInteraction, Id: "itx_b2"}, false, ""},
	}
	for _, c := range cases {
		sid, ok, err := g.Check(ctx, "alice", c.route)
		if err != nil {
			t.Errorf("%s: err %v", c.name, err)
			continue
		}
		if ok != c.ok || sid != c.sid {
			t.Errorf("%s: got ok=%v sid=%q, want ok=%v sid=%q", c.name, ok, sid, c.ok, c.sid)
		}
	}
	// 反查成功的三项应补写索引
	want := map[string]string{"run_a2": "ses_a1", "file_a2": "ses_a1", "itx_a2": "ses_a1"}
	if len(sink.got) != len(want) {
		t.Fatalf("补写索引 %v", sink.got)
	}
	for _, r := range sink.got {
		if want[r.Id] != r.SessionId {
			t.Errorf("补写了不应写入的索引 %+v", r)
		}
	}
}

func TestGuardStoreError(t *testing.T) {
	g := &Guard{Store: &fakeStore{err: errors.New("mongo down")}}
	if _, ok, err := g.Check(context.Background(), "alice", Route{Owner: OwnerSession, Id: "ses_a1"}); err == nil || ok {
		t.Fatalf("存储出错时应返回错误且不放行，得到 ok=%v err=%v", ok, err)
	}
}