Раскладка файлов после распознавния

This commit is contained in:
2026-06-14 14:53:40 +03:00
parent 91c501624a
commit 9c1b178e46
19 changed files with 3001 additions and 38 deletions
+550
View File
@@ -0,0 +1,550 @@
package worker
import (
"context"
"encoding/json"
"io"
"log/slog"
"os"
"path/filepath"
"testing"
"git.vakhrushev.me/av/jellybit/internal/layout"
"git.vakhrushev.me/av/jellybit/internal/qbt"
"git.vakhrushev.me/av/jellybit/internal/recognize"
"git.vakhrushev.me/av/jellybit/internal/store"
)
// memStore — полноценный in-memory store для тестов Ф3.
type memStore struct {
downloads map[int64]*store.Download
recs []*store.Recognition
hints map[int64][]string
overrides map[int64]map[string]string
links []store.FileLink
}
func newMemStore() *memStore {
return &memStore{
downloads: map[int64]*store.Download{},
hints: map[int64][]string{},
overrides: map[int64]map[string]string{},
}
}
func (m *memStore) put(d *store.Download) { m.downloads[d.ID] = d }
func (m *memStore) ListDownloadsByState(_ context.Context, states ...store.State) ([]store.Download, error) {
var out []store.Download
for _, d := range m.downloads {
for _, s := range states {
if d.State == s {
out = append(out, *d)
}
}
}
return out, nil
}
func (m *memStore) GetDownload(_ context.Context, id int64) (*store.Download, error) {
d, ok := m.downloads[id]
if !ok {
return nil, os.ErrNotExist
}
cp := *d
return &cp, nil
}
func (m *memStore) SetDownloadState(_ context.Context, id int64, st store.State, code, msg string) error {
d := m.downloads[id]
d.State = st
d.ErrorCode = store.NullString(code)
d.ErrorMsg = store.NullString(msg)
return nil
}
func (m *memStore) CreateRecognition(_ context.Context, r *store.Recognition, reasons []string) (int64, error) {
for _, e := range m.recs {
if e.DownloadID == r.DownloadID {
e.IsCurrent = false
}
}
cp := *r
cp.ID = int64(len(m.recs) + 1)
cp.IsCurrent = true
cp.AttemptNo = 1
for _, e := range m.recs {
if e.DownloadID == r.DownloadID {
cp.AttemptNo++
}
}
b, _ := jsonMarshal(reasons)
cp.Reasons = b
m.recs = append(m.recs, &cp)
return cp.ID, nil
}
func (m *memStore) GetCurrentRecognition(_ context.Context, downloadID int64) (*store.Recognition, error) {
for _, e := range m.recs {
if e.DownloadID == downloadID && e.IsCurrent {
cp := *e
return &cp, nil
}
}
return nil, nil
}
func (m *memStore) AddHint(_ context.Context, id int64, text string) error {
m.hints[id] = append(m.hints[id], text)
return nil
}
func (m *memStore) ListHints(_ context.Context, id int64) ([]string, error) { return m.hints[id], nil }
func (m *memStore) SetOverride(_ context.Context, id int64, field, value string) error {
if m.overrides[id] == nil {
m.overrides[id] = map[string]string{}
}
m.overrides[id][field] = value
return nil
}
func (m *memStore) ListOverrides(_ context.Context, id int64) (map[string]string, error) {
return m.overrides[id], nil
}
func (m *memStore) CreateFileLinks(_ context.Context, links []store.FileLink) error {
m.links = append(m.links, links...)
return nil
}
func (m *memStore) LatestBatchID(_ context.Context, id int64) (string, error) {
for i := len(m.links) - 1; i >= 0; i-- {
if m.links[i].DownloadID == id {
return m.links[i].ApplyBatchID, nil
}
}
return "", nil
}
func (m *memStore) ListFileLinksByBatch(_ context.Context, batch string) ([]store.FileLink, error) {
var out []store.FileLink
for _, l := range m.links {
if l.ApplyBatchID == batch {
out = append(out, l)
}
}
return out, nil
}
func (m *memStore) DeleteFileLinksByBatch(_ context.Context, batch string) error {
kept := m.links[:0]
for _, l := range m.links {
if l.ApplyBatchID != batch {
kept = append(kept, l)
}
}
m.links = kept
return nil
}
func jsonMarshal(v any) (string, error) {
b, err := json.Marshal(v)
return string(b), err
}
// fakeRecognizer возвращает заданный результат; onCall — побочный эффект для
// симуляции гонок (напр. отмена во время вызова LLM).
type fakeRecognizer struct {
result recognize.Result
err error
onCall func()
calls int
}
func (f *fakeRecognizer) Recognize(_ context.Context, _ recognize.Input) (recognize.Result, error) {
f.calls++
if f.onCall != nil {
f.onCall()
}
return f.result, f.err
}
func testWorkerWith(st Store, qb QBittorrent, rec Recognizer, lay Layouter) *Worker {
w := New(st, qb, rec, lay, Config{Category: "jellybit"},
slog.New(slog.NewTextHandler(io.Discard, nil)))
n := 0
w.newID = func() string { n++; return "batch-" + itoa(n) }
return w
}
func itoa(n int) string {
if n == 0 {
return "0"
}
var b []byte
for n > 0 {
b = append([]byte{byte('0' + n%10)}, b...)
n /= 10
}
return string(b)
}
const ihTest = "541adcff3b6dd5dba7088ea83317d9d6fac331d6"
func completedDownload(id int64) *store.Download {
return &store.Download{
ID: id, State: store.StateCompleted, SourceType: store.SourceMagnet,
SourceRef: "magnet:?xt=urn:btih:" + ihTest, Infohash: store.NullString(ihTest),
Context: "ctx",
}
}
func seriesResult() recognize.Result {
s, e1, e2 := 2, 1, 2
return recognize.Result{
Plan: recognize.Plan{
Type: recognize.MediaSeries, Title: "Show", Year: 2006, Confidence: 0.7,
Files: []recognize.PlanFile{
{Src: "Show/e1.mkv", Role: recognize.RoleEpisode, Season: &s, Episode: &e1},
{Src: "Show/e2.mkv", Role: recognize.RoleEpisode, Season: &s, Episode: &e2},
},
},
Decision: recognize.Decision{Reasons: []string{"нет матча в базе"}},
Raw: `{"type":"series"}`,
}
}
func TestRecognizeOne_CompletedToReview(t *testing.T) {
st := newMemStore()
st.put(completedDownload(1))
qb := &fakeQbt{
torrents: []qbt.Torrent{{Hash: ihTest, Name: "Show", SavePath: "/d", Category: "jellybit"}},
files: []qbt.File{{Name: "Show/e1.mkv", Size: 100}, {Name: "Show/e2.mkv", Size: 100}},
}
rec := &fakeRecognizer{result: seriesResult()}
w := testWorkerWith(st, qb, rec, nil)
w.recognizeOne(context.Background(), 1)
if st.downloads[1].State != store.StateReview {
t.Fatalf("state = %q, want review", st.downloads[1].State)
}
cur, _ := st.GetCurrentRecognition(context.Background(), 1)
if cur == nil || cur.Title.String != "Show" {
t.Fatalf("recognition = %+v", cur)
}
if !cur.Plan.Valid {
t.Error("plan must be persisted")
}
}
func TestRecognizeOne_DiscardsWhenStateChanged(t *testing.T) {
st := newMemStore()
st.put(completedDownload(1))
qb := &fakeQbt{
torrents: []qbt.Torrent{{Hash: ihTest, Name: "Show", SavePath: "/d"}},
files: []qbt.File{{Name: "Show/e1.mkv", Size: 100}},
}
// Во время вызова LLM задачу отменяют.
rec := &fakeRecognizer{result: seriesResult(), onCall: func() {
st.downloads[1].State = store.StateCancelled
}}
w := testWorkerWith(st, qb, rec, nil)
w.recognizeOne(context.Background(), 1)
if st.downloads[1].State != store.StateCancelled {
t.Errorf("state = %q, want cancelled (result discarded)", st.downloads[1].State)
}
if cur, _ := st.GetCurrentRecognition(context.Background(), 1); cur != nil {
t.Error("recognition must not be persisted after discard")
}
}
func TestRecognizeOne_SignalsErrorToReview(t *testing.T) {
st := newMemStore()
st.put(completedDownload(1))
qb := &fakeQbt{torrents: nil} // торрент пропал
rec := &fakeRecognizer{result: seriesResult()}
w := testWorkerWith(st, qb, rec, nil)
w.recognizeOne(context.Background(), 1)
if st.downloads[1].State != store.StateReview {
t.Fatalf("state = %q, want review", st.downloads[1].State)
}
cur, _ := st.GetCurrentRecognition(context.Background(), 1)
if cur == nil || len(cur.ReasonList()) == 0 {
t.Fatal("expected review with reason")
}
}
func TestRefine_AddsHintAndRerecognizes(t *testing.T) {
st := newMemStore()
d := completedDownload(1)
d.State = store.StateReview
st.put(d)
w := testWorkerWith(st, &fakeQbt{}, &fakeRecognizer{}, nil)
if err := w.Refine(context.Background(), 1, "это второй сезон"); err != nil {
t.Fatalf("Refine: %v", err)
}
if st.downloads[1].State != store.StateRecognizing {
t.Errorf("state = %q, want recognizing", st.downloads[1].State)
}
if h := st.hints[1]; len(h) != 1 || h[0] != "это второй сезон" {
t.Errorf("hints = %v", h)
}
if err := w.Refine(context.Background(), 1, " "); err == nil {
t.Error("empty hint must be rejected")
}
}
func TestSetType(t *testing.T) {
st := newMemStore()
d := completedDownload(1)
d.State = store.StateReview
st.put(d)
w := testWorkerWith(st, &fakeQbt{}, &fakeRecognizer{}, nil)
if err := w.SetType(context.Background(), 1, "series"); err != nil {
t.Fatalf("SetType: %v", err)
}
if st.overrides[1][ovrMediaType] != "series" {
t.Errorf("override = %v", st.overrides[1])
}
if st.downloads[1].State != store.StateRecognizing {
t.Errorf("state = %q, want recognizing", st.downloads[1].State)
}
if err := w.SetType(context.Background(), 1, "cartoon"); err == nil {
t.Error("invalid type must be rejected")
}
}
func TestIgnoreFile(t *testing.T) {
st := newMemStore()
d := completedDownload(1)
d.State = store.StateReview
st.put(d)
w := testWorkerWith(st, &fakeQbt{}, &fakeRecognizer{}, nil)
if err := w.IgnoreFile(context.Background(), 1, "Show/sample.mkv"); err != nil {
t.Fatalf("IgnoreFile: %v", err)
}
if err := w.IgnoreFile(context.Background(), 1, "Show/sample.mkv"); err != nil { // повтор не дублирует
t.Fatalf("IgnoreFile repeat: %v", err)
}
ignored := parseIgnored(st.overrides[1][ovrIgnoredFiles])
if len(ignored) != 1 || ignored[0] != "Show/sample.mkv" {
t.Errorf("ignored = %v", ignored)
}
if st.downloads[1].State != store.StateReview {
t.Errorf("ignore must keep review, got %q", st.downloads[1].State)
}
}
func TestDefer(t *testing.T) {
st := newMemStore()
d := completedDownload(1)
d.State = store.StateReview
st.put(d)
w := testWorkerWith(st, &fakeQbt{}, &fakeRecognizer{}, nil)
if err := w.Defer(context.Background(), 1); err != nil {
t.Fatalf("Defer: %v", err)
}
if st.downloads[1].State != store.StateDeferred {
t.Errorf("state = %q, want deferred", st.downloads[1].State)
}
}
// applyFixture — реальный layouter с temp-библиотеками и исходными файлами.
type applyFixture struct {
w *Worker
st *memStore
downloads string
movies string
series string
}
// newApplyFixture готовит worker с реальным layouter: исходные файлы лежат в
// downloads (он же savePath торрента), библиотеки — movies/series.
func newApplyFixture(t *testing.T, plan recognize.Plan) applyFixture {
t.Helper()
root := t.TempDir()
downloads := filepath.Join(root, "downloads")
movies := filepath.Join(root, "movies")
series := filepath.Join(root, "series")
for _, d := range []string{downloads, movies, series} {
_ = os.MkdirAll(d, 0o755)
}
for _, f := range plan.Files {
p := filepath.Join(downloads, f.Src)
_ = os.MkdirAll(filepath.Dir(p), 0o755)
if err := os.WriteFile(p, []byte("data-"+f.Src), 0o644); err != nil {
t.Fatal(err)
}
}
lay, err := layout.New(layout.Config{MoviesDir: movies, SeriesDir: series})
if err != nil {
t.Fatal(err)
}
st := newMemStore()
d := completedDownload(1)
d.State = store.StateReview
st.put(d)
planJSON, _ := json.Marshal(plan)
st.recs = append(st.recs, &store.Recognition{
ID: 1, DownloadID: 1, IsCurrent: true, Plan: store.NullString(string(planJSON)),
})
qb := &fakeQbt{torrents: []qbt.Torrent{{Hash: ihTest, SavePath: downloads, Category: "jellybit"}}}
w := testWorkerWith(st, qb, &fakeRecognizer{}, lay)
return applyFixture{w: w, st: st, downloads: downloads, movies: movies, series: series}
}
func TestApply_LinksAndDone(t *testing.T) {
f := newApplyFixture(t, seriesResult().Plan)
if err := f.w.Apply(context.Background(), 1); err != nil {
t.Fatalf("Apply: %v", err)
}
if f.st.downloads[1].State != store.StateDone {
t.Fatalf("state = %q, want done", f.st.downloads[1].State)
}
if len(f.st.links) != 2 {
t.Fatalf("file_links = %d, want 2", len(f.st.links))
}
want := filepath.Join(f.series, "Show (2006)", "Season 02", "Show (2006) S02E01.mkv")
if _, err := os.Stat(want); err != nil {
t.Errorf("expected hardlink %q: %v", want, err)
}
if _, err := os.Stat(filepath.Join(f.downloads, "Show/e1.mkv")); err != nil {
t.Errorf("source must remain: %v", err)
}
}
func TestApply_IgnoredFileSkipped(t *testing.T) {
plan := seriesResult().Plan
s, e := 2, 9
plan.Files = append(plan.Files, recognize.PlanFile{
Src: "Show/sample.mkv", Role: recognize.RoleEpisode, Season: &s, Episode: &e,
})
f := newApplyFixture(t, plan)
_ = f.st.SetOverride(context.Background(), 1, ovrIgnoredFiles, `["Show/sample.mkv"]`)
if err := f.w.Apply(context.Background(), 1); err != nil {
t.Fatalf("Apply: %v", err)
}
if len(f.st.links) != 2 { // sample пропущен
t.Errorf("file_links = %d, want 2 (sample ignored)", len(f.st.links))
}
}
func TestApply_CollisionStaysReview(t *testing.T) {
plan := seriesResult().Plan
f := newApplyFixture(t, plan)
// Занимаем цель первой серии чужим файлом.
dst := filepath.Join(f.series, "Show (2006)", "Season 02", "Show (2006) S02E01.mkv")
_ = os.MkdirAll(filepath.Dir(dst), 0o755)
_ = os.WriteFile(dst, []byte("foreign"), 0o644)
err := f.w.Apply(context.Background(), 1)
if err == nil {
t.Fatal("want collision error")
}
if f.st.downloads[1].State != store.StateReview {
t.Errorf("state = %q, want review after collision", f.st.downloads[1].State)
}
b, _ := os.ReadFile(dst)
if string(b) != "foreign" {
t.Errorf("foreign file overwritten: %q", b)
}
}
func TestUndo_RevertsLinks(t *testing.T) {
plan := seriesResult().Plan
f := newApplyFixture(t, plan)
if err := f.w.Apply(context.Background(), 1); err != nil {
t.Fatalf("Apply: %v", err)
}
dst := filepath.Join(f.series, "Show (2006)", "Season 02", "Show (2006) S02E01.mkv")
if _, err := os.Stat(dst); err != nil {
t.Fatalf("precondition: link must exist: %v", err)
}
if err := f.w.Undo(context.Background(), 1); err != nil {
t.Fatalf("Undo: %v", err)
}
if f.st.downloads[1].State != store.StateReverted {
t.Errorf("state = %q, want reverted", f.st.downloads[1].State)
}
if _, err := os.Stat(dst); !os.IsNotExist(err) {
t.Errorf("link must be removed: %v", err)
}
if len(f.st.links) != 0 {
t.Errorf("file_links must be deleted, got %d", len(f.st.links))
}
// Источник цел.
if _, err := os.Stat(filepath.Join(f.downloads, "Show/e1.mkv")); err != nil {
t.Errorf("source removed by undo: %v", err)
}
}
func TestReviewData(t *testing.T) {
plan := seriesResult().Plan
f := newApplyFixture(t, plan)
_ = f.st.AddHint(context.Background(), 1, "подсказка")
rd, err := f.w.ReviewData(context.Background(), 1)
if err != nil {
t.Fatalf("ReviewData: %v", err)
}
if rd.Recognition == nil || len(rd.Plan.Files) != 2 {
t.Fatalf("plan files = %+v", rd.Plan)
}
if len(rd.Preview) != 2 {
t.Errorf("preview links = %d, want 2", len(rd.Preview))
}
if len(rd.Hints) != 1 {
t.Errorf("hints = %v", rd.Hints)
}
}
func TestApplyOverrides(t *testing.T) {
plan := recognize.Plan{
Type: recognize.MediaMovie,
Files: []recognize.PlanFile{
{Src: "a.mkv", Role: recognize.RoleMain},
{Src: "b.mkv", Role: recognize.RoleEpisode},
},
}
out := applyOverrides(plan, map[string]string{
ovrMediaType: "series",
ovrIgnoredFiles: `["a.mkv"]`,
})
if out.Type != recognize.MediaSeries {
t.Errorf("type = %q, want series", out.Type)
}
if out.Files[0].Role != "ignore" {
t.Errorf("a.mkv role = %q, want ignore", out.Files[0].Role)
}
}
func TestToLayoutPlan(t *testing.T) {
s, e := 1, 3
plan := recognize.Plan{
Type: recognize.MediaSeries, Title: "X", Year: 2020,
Files: []recognize.PlanFile{
{Src: "e.mkv", Role: recognize.RoleEpisode, Season: &s, Episode: &e},
{Src: "sample.mkv", Role: "sample"},
},
}
lp := toLayoutPlan(plan, "/d")
if len(lp.Files) != 1 {
t.Fatalf("want 1 linkable file, got %d", len(lp.Files))
}
if lp.Files[0].Src != filepath.Join("/d", "e.mkv") {
t.Errorf("src = %q", lp.Files[0].Src)
}
if lp.Files[0].Role != layout.RoleEpisode {
t.Errorf("role = %q", lp.Files[0].Role)
}
}