transmission-manager/scheduler_test.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"
)

type fakeTransmission struct {
	mu        sync.Mutex
	torrents  map[string]schedulerTorrent
	failStart bool
	actions   []string
}

func newFakeTransmission(t *testing.T, torrents ...schedulerTorrent) (*Client, *fakeTransmission) {
	t.Helper()
	fake := &fakeTransmission{torrents: make(map[string]schedulerTorrent)}
	for _, torrent := range torrents {
		fake.torrents[torrent.Hash] = torrent
	}
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get(sessionHeader) != "test-session" {
			w.Header().Set(sessionHeader, "test-session")
			w.WriteHeader(http.StatusConflict)
			return
		}
		var request struct {
			Method    string `json:"method"`
			Arguments struct {
				IDs []string `json:"ids"`
			} `json:"arguments"`
		}
		if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
			t.Errorf("decode RPC: %v", err)
			return
		}
		fake.mu.Lock()
		defer fake.mu.Unlock()
		response := map[string]any{"result": "success", "arguments": map[string]any{}}
		switch request.Method {
		case "torrent-get":
			var result []schedulerTorrent
			for _, torrent := range fake.torrents {
				result = append(result, torrent)
			}
			response["arguments"] = map[string]any{"torrents": result}
		case "torrent-start-now":
			if fake.failStart {
				response["result"] = "start failed"
			} else {
				for _, hash := range request.Arguments.IDs {
					fake.actions = append(fake.actions, request.Method+":"+hash)
					torrent := fake.torrents[hash]
					torrent.Status = 6
					fake.torrents[hash] = torrent
				}
			}
		case "torrent-stop":
			for _, hash := range request.Arguments.IDs {
				fake.actions = append(fake.actions, request.Method+":"+hash)
				torrent := fake.torrents[hash]
				torrent.Status = 0
				fake.torrents[hash] = torrent
			}
		default:
			t.Errorf("unexpected RPC method %s", request.Method)
		}
		_ = json.NewEncoder(w).Encode(response)
	}))
	t.Cleanup(server.Close)
	return NewClient(server.URL, "", ""), fake
}

func testScheduler(t *testing.T, client *Client) *Scheduler {
	t.Helper()
	scheduler, err := NewScheduler(client, filepath.Join(t.TempDir(), "scheduler.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = scheduler.Close() })
	return scheduler
}

func addOptIn(t *testing.T, scheduler *Scheduler, torrents ...schedulerTorrent) {
	t.Helper()
	for _, torrent := range torrents {
		if _, err := scheduler.db.Exec(`INSERT INTO scheduler_torrents (hash, name) VALUES (?, ?)`, torrent.Hash, torrent.Name); err != nil {
			t.Fatal(err)
		}
	}
}

func activeHashes(t *testing.T, scheduler *Scheduler) []string {
	t.Helper()
	rows, err := scheduler.rows(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	var active []string
	for _, row := range rows {
		if row.Active {
			active = append(active, row.Hash)
		}
	}
	return active
}

func eventCount(t *testing.T, scheduler *Scheduler) int {
	t.Helper()
	var count int
	if err := scheduler.db.QueryRow("SELECT COUNT(*) FROM scheduler_events").Scan(&count); err != nil {
		t.Fatal(err)
	}
	return count
}

func TestSchedulerIdleRotationAndPersistence(t *testing.T) {
	a := schedulerTorrent{Hash: strings.Repeat("a", 40), Name: "A", PercentDone: 1}
	b := schedulerTorrent{Hash: strings.Repeat("b", 40), Name: "B", PercentDone: 1}
	c := schedulerTorrent{Hash: strings.Repeat("c", 40), Name: "C", PercentDone: 1}
	client, fake := newFakeTransmission(t, a, b, c)
	path := filepath.Join(t.TempDir(), "scheduler.db")
	scheduler, err := NewScheduler(client, path)
	if err != nil {
		t.Fatal(err)
	}
	addOptIn(t, scheduler, a, b, c)
	now := time.Unix(1_700_000_000, 0)
	scheduler.now = func() time.Time { return now }
	if err := scheduler.Tick(context.Background()); err != nil {
		t.Fatal(err)
	}
	if got := activeHashes(t, scheduler); len(got) != 1 || got[0] != a.Hash {
		t.Fatalf("initial active: %v", got)
	}
	now = now.Add(59 * time.Minute)
	if err := scheduler.Tick(context.Background()); err != nil {
		t.Fatal(err)
	}
	fake.mu.Lock()
	uploading := fake.torrents[a.Hash]
	uploading.UploadedEver = 10
	fake.torrents[a.Hash] = uploading
	fake.mu.Unlock()
	if err := scheduler.Tick(context.Background()); err != nil {
		t.Fatal(err)
	}
	now = now.Add(59 * time.Minute)
	if err := scheduler.Tick(context.Background()); err != nil {
		t.Fatal(err)
	}
	if got := activeHashes(t, scheduler); len(got) != 1 || got[0] != a.Hash {
		t.Fatalf("upload should delay rotation: %v", got)
	}
	now = now.Add(time.Minute)
	if err := scheduler.Tick(context.Background()); err != nil {
		t.Fatal(err)
	}
	if got := activeHashes(t, scheduler); len(got) != 1 || got[0] != b.Hash {
		t.Fatalf("expected oldest waiting torrent: %v", got)
	}
	if got := eventCount(t, scheduler); got != 2 {
		t.Fatalf("expected activation and rotation, got %d events", got)
	}
	if err := scheduler.Close(); err != nil {
		t.Fatal(err)
	}
	reopened, err := NewScheduler(client, path)
	if err != nil {
		t.Fatal(err)
	}
	defer reopened.Close()
	reopened.now = func() time.Time { return now }
	if err := reopened.Tick(context.Background()); err != nil {
		t.Fatal(err)
	}
	if got := activeHashes(t, reopened); len(got) != 1 || got[0] != b.Hash {
		t.Fatalf("restart lost slot state: %v", got)
	}
	if got := eventCount(t, reopened); got != 2 {
		t.Fatalf("restart lost history: %d", got)
	}
}

func TestSchedulerMultipleSlotsAndManualActions(t *testing.T) {
	a := schedulerTorrent{Hash: strings.Repeat("a", 40), Name: "A", PercentDone: 1}
	b := schedulerTorrent{Hash: strings.Repeat("b", 40), Name: "B", PercentDone: 1}
	c := schedulerTorrent{Hash: strings.Repeat("c", 40), Name: "C", PercentDone: 1}
	unmanagedDownload := schedulerTorrent{Hash: strings.Repeat("d", 40), Name: "Unmanaged download", PercentDone: .5, Status: 4}
	unmanagedSeed := schedulerTorrent{Hash: strings.Repeat("e", 40), Name: "Unmanaged seed", PercentDone: 1, Status: 6}
	client, fake := newFakeTransmission(t, a, b, c, unmanagedDownload, unmanagedSeed)
	scheduler := testScheduler(t, client)
	addOptIn(t, scheduler, a, b, c)
	if _, err := scheduler.db.Exec("UPDATE scheduler_settings SET max_active=2 WHERE id=1"); err != nil {
		t.Fatal(err)
	}
	now := time.Unix(1_700_000_000, 0)
	scheduler.now = func() time.Time { return now }
	if err := scheduler.Tick(context.Background()); err != nil {
		t.Fatal(err)
	}
	if got := activeHashes(t, scheduler); len(got) != 2 || got[0] != a.Hash || got[1] != b.Hash {
		t.Fatalf("expected two slots: %v", got)
	}
	fake.mu.Lock()
	paused := fake.torrents[a.Hash]
	paused.Status = 0
	fake.torrents[a.Hash] = paused
	resumed := fake.torrents[c.Hash]
	resumed.Status = 6
	fake.torrents[c.Hash] = resumed
	fake.mu.Unlock()
	if err := scheduler.Tick(context.Background()); err != nil {
		t.Fatal(err)
	}
	fake.mu.Lock()
	if fake.torrents[a.Hash].Status != 6 || fake.torrents[c.Hash].Status != 0 {
		t.Fatalf("manual actions were not reconciled: A=%d C=%d", fake.torrents[a.Hash].Status, fake.torrents[c.Hash].Status)
	}
	fake.mu.Unlock()
	now = now.Add(time.Hour)
	if err := scheduler.Tick(context.Background()); err != nil {
		t.Fatal(err)
	}
	if got := activeHashes(t, scheduler); len(got) != 2 || got[0] != b.Hash || got[1] != c.Hash {
		t.Fatalf("expected one rotation with no same-tick reuse: %v", got)
	}
	if _, err := scheduler.db.Exec("UPDATE scheduler_settings SET max_active=1 WHERE id=1"); err != nil {
		t.Fatal(err)
	}
	if err := scheduler.Tick(context.Background()); err != nil {
		t.Fatal(err)
	}
	if got := activeHashes(t, scheduler); len(got) != 1 {
		t.Fatalf("expected one opted-in active slot after reducing the limit: %v", got)
	}
	fake.mu.Lock()
	downloadStatus := fake.torrents[unmanagedDownload.Hash].Status
	seedStatus := fake.torrents[unmanagedSeed.Hash].Status
	actions := append([]string(nil), fake.actions...)
	fake.mu.Unlock()
	if downloadStatus != 4 || seedStatus != 6 {
		t.Fatalf("unopted torrents changed status: download=%d seed=%d", downloadStatus, seedStatus)
	}
	for _, action := range actions {
		if strings.Contains(action, unmanagedDownload.Hash) || strings.Contains(action, unmanagedSeed.Hash) {
			t.Fatalf("scheduler sent an action for an unopted torrent: %s", action)
		}
	}
}

func TestSchedulerFailureDoesNotRecordActivation(t *testing.T) {
	a := schedulerTorrent{Hash: strings.Repeat("a", 40), Name: "A", PercentDone: 1}
	client, fake := newFakeTransmission(t, a)
	scheduler := testScheduler(t, client)
	addOptIn(t, scheduler, a)
	fake.failStart = true
	if err := scheduler.Tick(context.Background()); err == nil {
		t.Fatal("expected RPC start failure")
	}
	if got := eventCount(t, scheduler); got != 0 {
		t.Fatalf("failed activation recorded %d events", got)
	}
	fake.failStart = false
	if err := scheduler.Tick(context.Background()); err != nil {
		t.Fatal(err)
	}
	if got := eventCount(t, scheduler); got != 1 {
		t.Fatalf("successful activation recorded %d events", got)
	}
}

func TestSchedulerKeepsOnlyTorrentWhenIdle(t *testing.T) {
	a := schedulerTorrent{Hash: strings.Repeat("a", 40), Name: "A", PercentDone: 1}
	client, _ := newFakeTransmission(t, a)
	scheduler := testScheduler(t, client)
	addOptIn(t, scheduler, a)
	now := time.Unix(1_700_000_000, 0)
	scheduler.now = func() time.Time { return now }
	if err := scheduler.Tick(context.Background()); err != nil {
		t.Fatal(err)
	}
	now = now.Add(2 * time.Hour)
	if err := scheduler.Tick(context.Background()); err != nil {
		t.Fatal(err)
	}
	if got := activeHashes(t, scheduler); len(got) != 1 || got[0] != a.Hash {
		t.Fatalf("only torrent should keep its slot: %v", got)
	}
	if got := eventCount(t, scheduler); got != 1 {
		t.Fatalf("idle slot without a candidate should not rotate: %d events", got)
	}
}

func TestSchedulerHTTPValidationAndHistory(t *testing.T) {
	a := schedulerTorrent{Hash: strings.Repeat("a", 40), Name: "A", PercentDone: .5}
	client, fake := newFakeTransmission(t, a)
	scheduler := testScheduler(t, client)
	request := func(handler http.HandlerFunc, method, path, body string) *httptest.ResponseRecorder {
		recorder := httptest.NewRecorder()
		handler(recorder, httptest.NewRequest(method, path, bytes.NewBufferString(body)))
		return recorder
	}
	if got := request(scheduler.handleSettings, "POST", "/api/scheduler/settings", `{"maxActive":101}`).Code; got != 400 {
		t.Fatalf("settings validation returned %d", got)
	}
	if got := request(scheduler.handleOptIn, "POST", "/api/scheduler/torrents/"+a.Hash+"/opt-in", `{"enabled":true}`).Code; got != 409 {
		t.Fatalf("incomplete opt-in returned %d", got)
	}
	if got := request(scheduler.handleSettings, "POST", "/api/scheduler/settings", `{"maxActive":2}`).Code; got != 200 {
		t.Fatalf("valid settings returned %d", got)
	}
	fake.mu.Lock()
	completed := fake.torrents[a.Hash]
	completed.PercentDone = 1
	fake.torrents[a.Hash] = completed
	fake.mu.Unlock()
	if got := request(scheduler.handleOptIn, "POST", "/api/scheduler/torrents/"+a.Hash+"/opt-in", `{"enabled":true}`).Code; got != 200 {
		t.Fatalf("valid opt-in returned %d", got)
	}
	if got := activeHashes(t, scheduler); len(got) != 1 || got[0] != a.Hash {
		t.Fatalf("opt-in did not fill a slot: %v", got)
	}
	if got := request(scheduler.handleOptIn, "POST", "/api/scheduler/torrents/"+a.Hash+"/opt-in", `{"enabled":false}`).Code; got != 200 {
		t.Fatalf("opt-out returned %d", got)
	}
	if got := activeHashes(t, scheduler); len(got) != 0 {
		t.Fatalf("opt-out did not release its slot: %v", got)
	}
	if _, err := scheduler.db.Exec("DELETE FROM scheduler_events"); err != nil {
		t.Fatal(err)
	}
	for i := 0; i < 51; i++ {
		if _, err := scheduler.db.Exec(`INSERT INTO scheduler_events
			(happened_at, from_hash, from_name, to_hash, to_name, reason) VALUES (1,'','','a','A','test')`); err != nil {
			t.Fatal(err)
		}
	}
	response := request(scheduler.handleHistory, "GET", "/api/scheduler/history", "")
	if response.Code != 200 {
		t.Fatalf("history returned %d: %s", response.Code, response.Body.String())
	}
	var page struct {
		Events  []schedulerEvent `json:"events"`
		HasMore bool             `json:"hasMore"`
	}
	if err := json.Unmarshal(response.Body.Bytes(), &page); err != nil {
		t.Fatal(err)
	}
	if len(page.Events) != 50 || !page.HasMore {
		t.Fatalf("bad first history page: %d events, more=%v", len(page.Events), page.HasMore)
	}
	response = request(scheduler.handleHistory, "GET", "/api/scheduler/history?before="+strconv.FormatInt(page.Events[49].ID, 10), "")
	if response.Code != 200 {
		t.Fatalf("history cursor returned %d", response.Code)
	}
	if err := json.Unmarshal(response.Body.Bytes(), &page); err != nil {
		t.Fatal(err)
	}
	if len(page.Events) != 1 || page.HasMore {
		t.Fatalf("bad final history page: %d events, more=%v", len(page.Events), page.HasMore)
	}
}