Раскладка файлов после распознавния
This commit is contained in:
@@ -0,0 +1,537 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"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"
|
||||
)
|
||||
|
||||
// Поля override.
|
||||
const (
|
||||
ovrMediaType = "media_type"
|
||||
ovrIgnoredFiles = "ignored_files"
|
||||
)
|
||||
|
||||
// recognizePending распознаёт завершённые загрузки и перезапускает те, что
|
||||
// помечены к перераспознаванию (recognizing — например, после подсказки или
|
||||
// после рестарта сервиса). Выполняется последовательно в поллинг-горутине;
|
||||
// сам вызов LLM идёт вне блокировки, поэтому команды ревью не простаивают.
|
||||
func (w *Worker) recognizePending(ctx context.Context) {
|
||||
w.mu.Lock()
|
||||
pending, err := w.store.ListDownloadsByState(ctx, store.StateCompleted, store.StateRecognizing)
|
||||
w.mu.Unlock()
|
||||
if err != nil {
|
||||
w.log.Warn("recognize: list pending failed", "err", err)
|
||||
return
|
||||
}
|
||||
for _, d := range pending {
|
||||
w.recognizeOne(ctx, d.ID)
|
||||
}
|
||||
}
|
||||
|
||||
// recognizeOne проводит одну загрузку через распознавание. Claim-паттерн:
|
||||
// под блокировкой переводим в recognizing, LLM зовём без блокировки, затем
|
||||
// под блокировкой фиксируем результат — но только если задачу за это время
|
||||
// не увели в другое состояние (cancel/defer).
|
||||
func (w *Worker) recognizeOne(ctx context.Context, id int64) {
|
||||
w.mu.Lock()
|
||||
d, err := w.store.GetDownload(ctx, id)
|
||||
if err != nil {
|
||||
w.mu.Unlock()
|
||||
w.log.Warn("recognize: get download", "download_id", id, "err", err)
|
||||
return
|
||||
}
|
||||
if d.State != store.StateCompleted && d.State != store.StateRecognizing {
|
||||
w.mu.Unlock()
|
||||
return
|
||||
}
|
||||
if d.State == store.StateCompleted {
|
||||
w.transition(ctx, *d, store.StateRecognizing, "", "")
|
||||
}
|
||||
w.mu.Unlock()
|
||||
|
||||
result, savePath, err := w.runRecognize(ctx, *d)
|
||||
if err != nil {
|
||||
// Не смогли получить сигналы или вызвать LLM — уходим в review с
|
||||
// причиной, человек перезапустит подсказкой.
|
||||
result = recognize.Result{Decision: recognize.Decision{
|
||||
Reasons: []string{"распознавание не удалось: " + err.Error()},
|
||||
}}
|
||||
}
|
||||
w.finishRecognition(ctx, id, result, savePath)
|
||||
}
|
||||
|
||||
// runRecognize собирает сигналы из qBittorrent и накопленные подсказки,
|
||||
// затем зовёт распознаватель. Возвращает также savePath для маппинга
|
||||
// относительных путей файлов в абсолютные при раскладке.
|
||||
func (w *Worker) runRecognize(ctx context.Context, d store.Download) (recognize.Result, string, error) {
|
||||
if !d.Infohash.Valid {
|
||||
return recognize.Result{}, "", fmt.Errorf("нет infohash")
|
||||
}
|
||||
t, ok, err := w.torrentByInfohash(ctx, d.Infohash.String)
|
||||
if err != nil {
|
||||
return recognize.Result{}, "", err
|
||||
}
|
||||
if !ok {
|
||||
return recognize.Result{}, "", fmt.Errorf("торрент не найден в qBittorrent")
|
||||
}
|
||||
files, err := w.qbt.Files(ctx, t.Hash)
|
||||
if err != nil {
|
||||
return recognize.Result{}, "", err
|
||||
}
|
||||
hints, err := w.store.ListHints(ctx, d.ID)
|
||||
if err != nil {
|
||||
return recognize.Result{}, "", err
|
||||
}
|
||||
|
||||
in := recognize.Input{
|
||||
Name: t.Name,
|
||||
Context: d.Context,
|
||||
Hints: hints,
|
||||
Files: make([]recognize.File, len(files)),
|
||||
}
|
||||
for i, f := range files {
|
||||
in.Files[i] = recognize.File{Path: f.Name, Size: f.Size}
|
||||
}
|
||||
|
||||
res, err := w.recognizer.Recognize(ctx, in)
|
||||
if err != nil {
|
||||
return recognize.Result{}, t.SavePath, err
|
||||
}
|
||||
return res, t.SavePath, nil
|
||||
}
|
||||
|
||||
// finishRecognition сохраняет попытку распознавания и двигает задачу. В Ф3
|
||||
// метабазы выключены → авто-раскладки не делаем, всегда уходим в review.
|
||||
func (w *Worker) finishRecognition(ctx context.Context, id int64, res recognize.Result, _ string) {
|
||||
planJSON, err := json.Marshal(res.Plan)
|
||||
if err != nil {
|
||||
w.log.Error("recognize: marshal plan", "download_id", id, "err", err)
|
||||
planJSON = []byte("{}")
|
||||
}
|
||||
|
||||
rec := &store.Recognition{
|
||||
DownloadID: id,
|
||||
MediaType: store.NullString(string(res.Plan.Type)),
|
||||
Title: store.NullString(res.Plan.Title),
|
||||
Provider: store.NullString("none"),
|
||||
Plan: store.NullString(string(planJSON)),
|
||||
RawLLM: store.NullString(res.Raw),
|
||||
}
|
||||
if res.Plan.OriginalTitle != "" {
|
||||
rec.OriginalTitle = store.NullString(res.Plan.OriginalTitle)
|
||||
}
|
||||
if res.Plan.Year != 0 {
|
||||
rec.Year = sql.NullInt64{Int64: int64(res.Plan.Year), Valid: true}
|
||||
}
|
||||
if res.Plan.Confidence != 0 {
|
||||
rec.Confidence = sql.NullFloat64{Float64: res.Plan.Confidence, Valid: true}
|
||||
}
|
||||
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
|
||||
d, err := w.store.GetDownload(ctx, id)
|
||||
if err != nil {
|
||||
w.log.Warn("recognize: reload download", "download_id", id, "err", err)
|
||||
return
|
||||
}
|
||||
if d.State != store.StateRecognizing {
|
||||
// За время вызова LLM задачу увели (cancel/defer) — результат не нужен.
|
||||
w.log.Info("recognize: result discarded, state changed",
|
||||
"download_id", id, "state", d.State)
|
||||
return
|
||||
}
|
||||
if _, err := w.store.CreateRecognition(ctx, rec, res.Decision.Reasons); err != nil {
|
||||
w.log.Error("recognize: persist", "download_id", id, "err", err)
|
||||
return
|
||||
}
|
||||
w.transition(ctx, *d, store.StateReview, "", "")
|
||||
}
|
||||
|
||||
// --- Команды ревью ---
|
||||
|
||||
// Apply создаёт хардлинки по текущему плану (с применёнными правками) и
|
||||
// переводит задачу в done. Коллизия цели → остаёмся в review с причиной.
|
||||
func (w *Worker) Apply(ctx context.Context, id int64) error {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
if w.layouter == nil {
|
||||
return fmt.Errorf("apply: раскладчик не сконфигурирован")
|
||||
}
|
||||
|
||||
d, err := w.store.GetDownload(ctx, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("apply: %w", err)
|
||||
}
|
||||
if d.State != store.StateReview && d.State != store.StateDeferred {
|
||||
return fmt.Errorf("apply: задача %d в состоянии %s (ожидалось review/deferred)", id, d.State)
|
||||
}
|
||||
|
||||
plan, err := w.effectivePlan(ctx, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("apply: %w", err)
|
||||
}
|
||||
t, ok, err := w.torrentByInfohash(ctx, d.Infohash.String)
|
||||
if err != nil || !ok {
|
||||
return fmt.Errorf("apply: торрент не найден: %v", err)
|
||||
}
|
||||
|
||||
links, err := w.layouter.BuildLinks(toLayoutPlan(plan, t.SavePath))
|
||||
if err != nil {
|
||||
return fmt.Errorf("apply: построение ссылок: %w", err)
|
||||
}
|
||||
|
||||
batch := w.newID()
|
||||
results, applyErr := w.layouter.Apply(ctx, links)
|
||||
|
||||
// Фиксируем то, что успели слинковать (идемпотентность повторного apply).
|
||||
fl := make([]store.FileLink, 0, len(results))
|
||||
for _, r := range results {
|
||||
fl = append(fl, store.FileLink{
|
||||
DownloadID: id,
|
||||
ApplyBatchID: batch,
|
||||
SrcPath: r.Link.Src,
|
||||
DstPath: r.Link.Dst,
|
||||
Kind: string(r.Link.Kind),
|
||||
Status: string(r.Status),
|
||||
})
|
||||
}
|
||||
if len(fl) > 0 {
|
||||
if err := w.store.CreateFileLinks(ctx, fl); err != nil {
|
||||
return fmt.Errorf("apply: запись ссылок: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if applyErr != nil {
|
||||
if errors.Is(applyErr, layout.ErrCollision) {
|
||||
w.transition(ctx, *d, store.StateReview, "collision", applyErr.Error())
|
||||
return fmt.Errorf("apply: %w", applyErr)
|
||||
}
|
||||
w.transition(ctx, *d, store.StateFailed, "apply", applyErr.Error())
|
||||
return fmt.Errorf("apply: %w", applyErr)
|
||||
}
|
||||
|
||||
w.transition(ctx, *d, store.StateDone, "", "")
|
||||
w.log.Info("apply: linked", "download_id", id, "batch", batch, "links", len(fl))
|
||||
return nil
|
||||
}
|
||||
|
||||
// Refine добавляет подсказку и отправляет задачу на перераспознавание.
|
||||
func (w *Worker) Refine(ctx context.Context, id int64, hint string) error {
|
||||
hint = strings.TrimSpace(hint)
|
||||
if hint == "" {
|
||||
return fmt.Errorf("refine: пустая подсказка")
|
||||
}
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
|
||||
d, err := w.requireReviewable(ctx, id, "refine")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := w.store.AddHint(ctx, id, hint); err != nil {
|
||||
return fmt.Errorf("refine: %w", err)
|
||||
}
|
||||
w.transition(ctx, *d, store.StateRecognizing, "", "")
|
||||
return nil
|
||||
}
|
||||
|
||||
// SetType фиксирует тип (override) и перезапускает распознавание с подсказкой
|
||||
// — чтобы LLM пересобрал роли файлов под новый тип.
|
||||
func (w *Worker) SetType(ctx context.Context, id int64, mediaType string) error {
|
||||
if mediaType != string(recognize.MediaMovie) && mediaType != string(recognize.MediaSeries) {
|
||||
return fmt.Errorf("set type: недопустимый тип %q", mediaType)
|
||||
}
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
|
||||
d, err := w.requireReviewable(ctx, id, "set type")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := w.store.SetOverride(ctx, id, ovrMediaType, mediaType); err != nil {
|
||||
return fmt.Errorf("set type: %w", err)
|
||||
}
|
||||
label := "фильм"
|
||||
if mediaType == string(recognize.MediaSeries) {
|
||||
label = "сериал"
|
||||
}
|
||||
if err := w.store.AddHint(ctx, id, "Тип точно: "+label+"."); err != nil {
|
||||
return fmt.Errorf("set type: %w", err)
|
||||
}
|
||||
w.transition(ctx, *d, store.StateRecognizing, "", "")
|
||||
return nil
|
||||
}
|
||||
|
||||
// IgnoreFile помечает файл к игнорированию (не линкуем). Остаёмся в review;
|
||||
// превью пересчитается с учётом правки.
|
||||
func (w *Worker) IgnoreFile(ctx context.Context, id int64, src string) error {
|
||||
src = strings.TrimSpace(src)
|
||||
if src == "" {
|
||||
return fmt.Errorf("ignore: пустой путь")
|
||||
}
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
|
||||
if _, err := w.requireReviewable(ctx, id, "ignore"); err != nil {
|
||||
return err
|
||||
}
|
||||
overrides, err := w.store.ListOverrides(ctx, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("ignore: %w", err)
|
||||
}
|
||||
ignored := parseIgnored(overrides[ovrIgnoredFiles])
|
||||
if !contains(ignored, src) {
|
||||
ignored = append(ignored, src)
|
||||
}
|
||||
b, _ := json.Marshal(ignored)
|
||||
if err := w.store.SetOverride(ctx, id, ovrIgnoredFiles, string(b)); err != nil {
|
||||
return fmt.Errorf("ignore: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Defer паркует задачу в deferred (вернётся в ревью по действию).
|
||||
func (w *Worker) Defer(ctx context.Context, id int64) error {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
|
||||
d, err := w.store.GetDownload(ctx, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("defer: %w", err)
|
||||
}
|
||||
if d.State.IsTerminal() {
|
||||
return fmt.Errorf("defer: задача %d терминальна (%s)", id, d.State)
|
||||
}
|
||||
w.transition(ctx, *d, store.StateDeferred, "", "")
|
||||
return nil
|
||||
}
|
||||
|
||||
// Undo снимает хардлинки последнего батча и переводит задачу в reverted.
|
||||
// Источник недосягаем (раскладчик удаляет только пути под библиотекой).
|
||||
func (w *Worker) Undo(ctx context.Context, id int64) error {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
if w.layouter == nil {
|
||||
return fmt.Errorf("undo: раскладчик не сконфигурирован")
|
||||
}
|
||||
|
||||
d, err := w.store.GetDownload(ctx, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("undo: %w", err)
|
||||
}
|
||||
if d.State != store.StateDone {
|
||||
return fmt.Errorf("undo: задача %d в состоянии %s (ожидалось done)", id, d.State)
|
||||
}
|
||||
batch, err := w.store.LatestBatchID(ctx, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("undo: %w", err)
|
||||
}
|
||||
if batch == "" {
|
||||
return fmt.Errorf("undo: нечего откатывать")
|
||||
}
|
||||
rows, err := w.store.ListFileLinksByBatch(ctx, batch)
|
||||
if err != nil {
|
||||
return fmt.Errorf("undo: %w", err)
|
||||
}
|
||||
links := make([]layout.Link, len(rows))
|
||||
for i, r := range rows {
|
||||
links[i] = layout.Link{Src: r.SrcPath, Dst: r.DstPath, Kind: layout.Kind(r.Kind)}
|
||||
}
|
||||
n, err := w.layouter.Undo(ctx, links)
|
||||
if err != nil {
|
||||
return fmt.Errorf("undo: %w", err)
|
||||
}
|
||||
if err := w.store.DeleteFileLinksByBatch(ctx, batch); err != nil {
|
||||
return fmt.Errorf("undo: %w", err)
|
||||
}
|
||||
w.transition(ctx, *d, store.StateReverted, "", "")
|
||||
w.log.Info("undo: reverted", "download_id", id, "batch", batch, "removed", n)
|
||||
return nil
|
||||
}
|
||||
|
||||
// requireReviewable проверяет, что задача в review/deferred. Вызывается под mu.
|
||||
func (w *Worker) requireReviewable(ctx context.Context, id int64, op string) (*store.Download, error) {
|
||||
d, err := w.store.GetDownload(ctx, id)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%s: %w", op, err)
|
||||
}
|
||||
if d.State != store.StateReview && d.State != store.StateDeferred {
|
||||
return nil, fmt.Errorf("%s: задача %d в состоянии %s (ожидалось review/deferred)", op, id, d.State)
|
||||
}
|
||||
return d, nil
|
||||
}
|
||||
|
||||
// --- Данные для экрана ревью ---
|
||||
|
||||
// ReviewData — всё, что нужно транспорту для отрисовки ревью.
|
||||
type ReviewData struct {
|
||||
Download store.Download
|
||||
Recognition *store.Recognition
|
||||
Plan recognize.Plan // эффективный (с применёнными правками)
|
||||
Preview []layout.Link // целевые пути (Src — относительный, для показа)
|
||||
Hints []string
|
||||
Overrides map[string]string
|
||||
}
|
||||
|
||||
// ReviewData собирает данные ревью по загрузке.
|
||||
func (w *Worker) ReviewData(ctx context.Context, id int64) (*ReviewData, error) {
|
||||
d, err := w.store.GetDownload(ctx, id)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("review data: %w", err)
|
||||
}
|
||||
rec, err := w.store.GetCurrentRecognition(ctx, id)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("review data: %w", err)
|
||||
}
|
||||
hints, err := w.store.ListHints(ctx, id)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("review data: %w", err)
|
||||
}
|
||||
overrides, err := w.store.ListOverrides(ctx, id)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("review data: %w", err)
|
||||
}
|
||||
|
||||
rd := &ReviewData{Download: *d, Recognition: rec, Hints: hints, Overrides: overrides}
|
||||
if rec != nil && rec.Plan.Valid {
|
||||
var plan recognize.Plan
|
||||
if err := json.Unmarshal([]byte(rec.Plan.String), &plan); err == nil {
|
||||
plan = applyOverrides(plan, overrides)
|
||||
rd.Plan = plan
|
||||
// Превью строим по относительным путям; ошибку игнорируем —
|
||||
// просто покажем причины без превью.
|
||||
if w.layouter != nil {
|
||||
if links, lerr := w.layouter.BuildLinks(toLayoutPlan(plan, "")); lerr == nil {
|
||||
rd.Preview = links
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return rd, nil
|
||||
}
|
||||
|
||||
// effectivePlan загружает текущий план и применяет правки (под mu).
|
||||
func (w *Worker) effectivePlan(ctx context.Context, id int64) (recognize.Plan, error) {
|
||||
rec, err := w.store.GetCurrentRecognition(ctx, id)
|
||||
if err != nil {
|
||||
return recognize.Plan{}, err
|
||||
}
|
||||
if rec == nil || !rec.Plan.Valid {
|
||||
return recognize.Plan{}, fmt.Errorf("нет плана распознавания")
|
||||
}
|
||||
var plan recognize.Plan
|
||||
if err := json.Unmarshal([]byte(rec.Plan.String), &plan); err != nil {
|
||||
return recognize.Plan{}, fmt.Errorf("разбор плана: %w", err)
|
||||
}
|
||||
overrides, err := w.store.ListOverrides(ctx, id)
|
||||
if err != nil {
|
||||
return recognize.Plan{}, err
|
||||
}
|
||||
return applyOverrides(plan, overrides), nil
|
||||
}
|
||||
|
||||
// --- Хелперы преобразования ---
|
||||
|
||||
// applyOverrides применяет ручные правки к плану: форсит тип и помечает
|
||||
// игнорируемые файлы ролью ignore (их раскладка пропустит).
|
||||
func applyOverrides(plan recognize.Plan, overrides map[string]string) recognize.Plan {
|
||||
if mt := overrides[ovrMediaType]; mt == string(recognize.MediaMovie) || mt == string(recognize.MediaSeries) {
|
||||
plan.Type = recognize.MediaType(mt)
|
||||
}
|
||||
ignored := parseIgnored(overrides[ovrIgnoredFiles])
|
||||
if len(ignored) > 0 {
|
||||
for i := range plan.Files {
|
||||
if contains(ignored, plan.Files[i].Src) {
|
||||
plan.Files[i].Role = "ignore"
|
||||
}
|
||||
}
|
||||
}
|
||||
return plan
|
||||
}
|
||||
|
||||
// toLayoutPlan переводит план распознавания в план раскладки. srcPrefix
|
||||
// (savePath) приклеивается к относительным путям файлов; пустой — оставляет
|
||||
// относительные (для превью). Роли вне main/episode/subtitle отбрасываются.
|
||||
func toLayoutPlan(plan recognize.Plan, srcPrefix string) layout.Plan {
|
||||
lp := layout.Plan{
|
||||
Type: layout.MediaType(plan.Type),
|
||||
Title: plan.Title,
|
||||
Year: plan.Year,
|
||||
}
|
||||
for _, f := range plan.Files {
|
||||
role, ok := mapRole(f.Role)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
src := f.Src
|
||||
if srcPrefix != "" {
|
||||
src = filepath.Join(srcPrefix, f.Src)
|
||||
}
|
||||
lp.Files = append(lp.Files, layout.PlanFile{
|
||||
Src: src,
|
||||
Role: role,
|
||||
Season: f.Season,
|
||||
Episode: f.Episode,
|
||||
})
|
||||
}
|
||||
return lp
|
||||
}
|
||||
|
||||
func mapRole(r recognize.FileRole) (layout.Role, bool) {
|
||||
switch r {
|
||||
case recognize.RoleMain:
|
||||
return layout.RoleMain, true
|
||||
case recognize.RoleEpisode:
|
||||
return layout.RoleEpisode, true
|
||||
case recognize.RoleSubtitle:
|
||||
return layout.RoleSubtitle, true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
// torrentByInfohash ищет торрент категории по infohash (v1/v2/hash).
|
||||
func (w *Worker) torrentByInfohash(ctx context.Context, infohash string) (qbt.Torrent, bool, error) {
|
||||
torrents, err := w.qbt.Torrents(ctx, w.cfg.Category)
|
||||
if err != nil {
|
||||
return qbt.Torrent{}, false, err
|
||||
}
|
||||
want := strings.ToLower(infohash)
|
||||
for _, t := range torrents {
|
||||
for _, h := range []string{t.Hash, t.InfohashV1, t.InfohashV2} {
|
||||
if h != "" && strings.ToLower(h) == want {
|
||||
return t, true, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
return qbt.Torrent{}, false, nil
|
||||
}
|
||||
|
||||
func parseIgnored(s string) []string {
|
||||
if s == "" {
|
||||
return nil
|
||||
}
|
||||
var out []string
|
||||
_ = json.Unmarshal([]byte(s), &out)
|
||||
return out
|
||||
}
|
||||
|
||||
func contains(ss []string, s string) bool {
|
||||
for _, x := range ss {
|
||||
if x == s {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
+63
-10
@@ -4,7 +4,11 @@
|
||||
// состояние.
|
||||
//
|
||||
// Ф1 ведёт задачу downloading → completed, плюс stuck/failed по таймаутам и
|
||||
// ошибкам qBittorrent. Распознавание и раскладка (completed →) — Ф2+.
|
||||
// ошибкам qBittorrent. Ф3 продолжает: completed → recognizing (вызов
|
||||
// recognize) → review; команды ревью (apply/refine/reject/defer/undo,
|
||||
// переключение типа, пометка «игнор») раскладывают файлы хардлинками через
|
||||
// layout. Распознавание зовётся в поллинг-цикле, команды — из транспортов;
|
||||
// всё под per-download блокировкой w.mu.
|
||||
package worker
|
||||
|
||||
import (
|
||||
@@ -15,7 +19,9 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"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"
|
||||
)
|
||||
|
||||
@@ -24,12 +30,37 @@ type Store interface {
|
||||
ListDownloadsByState(ctx context.Context, states ...store.State) ([]store.Download, error)
|
||||
GetDownload(ctx context.Context, id int64) (*store.Download, error)
|
||||
SetDownloadState(ctx context.Context, id int64, state store.State, errCode, errMsg string) error
|
||||
|
||||
// Ф3: распознавание, ревью, раскладка.
|
||||
CreateRecognition(ctx context.Context, r *store.Recognition, reasons []string) (int64, error)
|
||||
GetCurrentRecognition(ctx context.Context, downloadID int64) (*store.Recognition, error)
|
||||
AddHint(ctx context.Context, downloadID int64, text string) error
|
||||
ListHints(ctx context.Context, downloadID int64) ([]string, error)
|
||||
SetOverride(ctx context.Context, downloadID int64, field, value string) error
|
||||
ListOverrides(ctx context.Context, downloadID int64) (map[string]string, error)
|
||||
CreateFileLinks(ctx context.Context, links []store.FileLink) error
|
||||
LatestBatchID(ctx context.Context, downloadID int64) (string, error)
|
||||
ListFileLinksByBatch(ctx context.Context, batchID string) ([]store.FileLink, error)
|
||||
DeleteFileLinksByBatch(ctx context.Context, batchID string) error
|
||||
}
|
||||
|
||||
// QBittorrent — нужная worker часть клиента qBittorrent.
|
||||
type QBittorrent interface {
|
||||
Torrents(ctx context.Context, category string) ([]qbt.Torrent, error)
|
||||
Add(ctx context.Context, ar qbt.AddRequest) error
|
||||
Files(ctx context.Context, hash string) ([]qbt.File, error)
|
||||
}
|
||||
|
||||
// Recognizer — распознаватель (recognize.Recognizer).
|
||||
type Recognizer interface {
|
||||
Recognize(ctx context.Context, in recognize.Input) (recognize.Result, error)
|
||||
}
|
||||
|
||||
// Layouter — раскладчик хардлинками (layout.Layouter).
|
||||
type Layouter interface {
|
||||
BuildLinks(p layout.Plan) ([]layout.Link, error)
|
||||
Apply(ctx context.Context, links []layout.Link) ([]layout.Result, error)
|
||||
Undo(ctx context.Context, links []layout.Link) (int, error)
|
||||
}
|
||||
|
||||
// Config — параметры воркера.
|
||||
@@ -43,18 +74,36 @@ type Config struct {
|
||||
|
||||
// Worker — поллер и владелец переходов.
|
||||
type Worker struct {
|
||||
store Store
|
||||
qbt QBittorrent
|
||||
cfg Config
|
||||
log *slog.Logger
|
||||
store Store
|
||||
qbt QBittorrent
|
||||
recognizer Recognizer
|
||||
layouter Layouter
|
||||
cfg Config
|
||||
log *slog.Logger
|
||||
|
||||
mu sync.Mutex // сериализует переходы (поллинг + команды)
|
||||
now func() time.Time // подменяется в тестах
|
||||
mu sync.Mutex // сериализует переходы (поллинг + команды)
|
||||
now func() time.Time // подменяется в тестах
|
||||
newID func() string // генератор apply_batch_id (подменяется в тестах)
|
||||
}
|
||||
|
||||
// New собирает воркер.
|
||||
func New(st Store, qb QBittorrent, cfg Config, log *slog.Logger) *Worker {
|
||||
return &Worker{store: st, qbt: qb, cfg: cfg, log: log, now: time.Now}
|
||||
// New собирает воркер. recognizer/layouter могут быть nil (Ф1 без Ф3-ступеней
|
||||
// распознавания и раскладки) — тогда completed-задачи не двигаются дальше.
|
||||
func New(st Store, qb QBittorrent, rec Recognizer, lay Layouter, cfg Config, log *slog.Logger) *Worker {
|
||||
return &Worker{
|
||||
store: st,
|
||||
qbt: qb,
|
||||
recognizer: rec,
|
||||
layouter: lay,
|
||||
cfg: cfg,
|
||||
log: log,
|
||||
now: time.Now,
|
||||
newID: defaultBatchID,
|
||||
}
|
||||
}
|
||||
|
||||
// defaultBatchID — уникальный идентификатор батча раскладки.
|
||||
func defaultBatchID() string {
|
||||
return fmt.Sprintf("b-%d", time.Now().UnixNano())
|
||||
}
|
||||
|
||||
// Run крутит цикл поллинга до отмены ctx.
|
||||
@@ -79,6 +128,10 @@ func (w *Worker) pollOnce(ctx context.Context) {
|
||||
if err := w.Poll(ctx); err != nil {
|
||||
w.log.Warn("poll failed", "err", err)
|
||||
}
|
||||
// Ф3: распознаём завершённые загрузки (и перезапускаем по подсказке).
|
||||
if w.recognizer != nil {
|
||||
w.recognizePending(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
// Poll сверяет активные задачи с состоянием qBittorrent и двигает их.
|
||||
|
||||
@@ -63,9 +63,31 @@ func (f *fakeStore) SetDownloadState(_ context.Context, id int64, st store.State
|
||||
return nil
|
||||
}
|
||||
|
||||
// --- Ф3-методы Store (заглушки; переопределяются в review_test.go) ---
|
||||
|
||||
func (f *fakeStore) CreateRecognition(_ context.Context, _ *store.Recognition, _ []string) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (f *fakeStore) GetCurrentRecognition(_ context.Context, _ int64) (*store.Recognition, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (f *fakeStore) AddHint(_ context.Context, _ int64, _ string) error { return nil }
|
||||
func (f *fakeStore) ListHints(_ context.Context, _ int64) ([]string, error) { return nil, nil }
|
||||
func (f *fakeStore) SetOverride(_ context.Context, _ int64, _, _ string) error { return nil }
|
||||
func (f *fakeStore) ListOverrides(_ context.Context, _ int64) (map[string]string, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (f *fakeStore) CreateFileLinks(_ context.Context, _ []store.FileLink) error { return nil }
|
||||
func (f *fakeStore) LatestBatchID(_ context.Context, _ int64) (string, error) { return "", nil }
|
||||
func (f *fakeStore) ListFileLinksByBatch(_ context.Context, _ string) ([]store.FileLink, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (f *fakeStore) DeleteFileLinksByBatch(_ context.Context, _ string) error { return nil }
|
||||
|
||||
type fakeQbt struct {
|
||||
torrents []qbt.Torrent
|
||||
added []qbt.AddRequest
|
||||
files []qbt.File
|
||||
}
|
||||
|
||||
func (f *fakeQbt) Torrents(_ context.Context, _ string) ([]qbt.Torrent, error) {
|
||||
@@ -77,8 +99,12 @@ func (f *fakeQbt) Add(_ context.Context, ar qbt.AddRequest) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeQbt) Files(_ context.Context, _ string) ([]qbt.File, error) {
|
||||
return f.files, nil
|
||||
}
|
||||
|
||||
func newTestWorker(st *fakeStore, qb *fakeQbt) *Worker {
|
||||
w := New(st, qb, Config{
|
||||
w := New(st, qb, nil, nil, Config{
|
||||
Category: "jellybit",
|
||||
SavePath: "/srv/media/downloads",
|
||||
MagnetTimeout: 30 * time.Minute,
|
||||
|
||||
Reference in New Issue
Block a user