eh-downloader/internal/eh/transport_test.go

package eh

// 连接复用测试
//
// 曾经的 bug：httpClient() 每次请求都新建 http.Transport，用完丢弃。
// Transport 被丢弃后，它持有的 keep-alive 连接不会立刻释放（要等
// IdleConnTimeout = 90s），于是每发一个请求就留下一条僵尸连接。
//
// 实测踩坑现场：解析大画廊时 91 条 Established 连接，而解析是串行的，
// 正常应该只有 1 条。连接堆积会吃代理的连接配额，也让新请求变慢。

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

// buildGalleryHTML 造一页画廊 HTML
func buildGalleryHTML(gid string, from, to, total, screen int) string {
	var sb strings.Builder
	if screen == 0 {
		sb.WriteString(`<!DOCTYPE html><html><head><title>[Test] - E-Hentai Galleries</title>`)
		sb.WriteString(`<meta name="description" content="Free Hentai Doujinshi Gallery: [Test]: x" />`)
		sb.WriteString(`</head><body><div id="taglist"><table></table></div>`)
	}
	for i := from; i <= to; i++ {
		sb.WriteString(fmt.Sprintf(
			`<a href="/s/k%03d/%s-%d"><div title="Page %d: %03d.png"></div></a>`,
			i, gid, i, i, i))
	}
	sb.WriteString(fmt.Sprintf(`Showing %d - %d of %d images`, from, to, total))
	if screen == 0 {
		sb.WriteString(`</body></html>`)
	}
	return sb.String()
}

// 统计服务端实际接受的连接数
func TestTransportIsReused(t *testing.T) {
	var conns int64

	srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("ok"))
	}))
	srv.Config.ConnState = func(c net.Conn, state http.ConnState) {
		if state == http.StateNew {
			atomic.AddInt64(&conns, 1)
		}
	}
	srv.Start()
	defer srv.Close()

	c := NewClient(10*time.Second, "", "off", "auto")

	const n = 20
	for i := 0; i < n; i++ {
		if _, err := c.getText(context.Background(), srv.URL+"/x"); err != nil {
			t.Fatalf("第 %d 次请求失败: %v", i+1, err)
		}
	}

	got := atomic.LoadInt64(&conns)
	// 连接复用后，20 次请求应该只用 1 条连接（串行）。
	// 修复前是每请求一条 = 20 条。
	if got > 3 {
		t.Errorf("%d 次请求开了 %d 条连接 —— Transport 没复用，连接在泄漏", n, got)
	}
	t.Logf("%d 次请求共用了 %d 条连接", n, got)
}

// 并发请求也应该复用连接，而不是每个 goroutine 一条
func TestTransportReusedUnderConcurrency(t *testing.T) {
	var conns int64

	srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("ok"))
	}))
	srv.Config.ConnState = func(c net.Conn, state http.ConnState) {
		if state == http.StateNew {
			atomic.AddInt64(&conns, 1)
		}
	}
	srv.Start()
	defer srv.Close()

	c := NewClient(10*time.Second, "", "off", "auto")

	var wg sync.WaitGroup
	for i := 0; i < 20; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for j := 0; j < 5; j++ {
				c.getText(context.Background(), srv.URL+"/x")
			}
		}()
	}
	wg.Wait()

	got := atomic.LoadInt64(&conns)
	// 100 次并发请求，连接数应该远小于 100
	if got > 60 {
		t.Errorf("100 次并发请求开了 %d 条连接，复用不充分", got)
	}
	t.Logf("100 次并发请求共用了 %d 条连接", got)
}

// Close() 要真的能关掉空闲连接
func TestClientCloseReleasesConnections(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("ok"))
	}))
	defer srv.Close()

	c := NewClient(10*time.Second, "", "off", "auto")
	for i := 0; i < 5; i++ {
		c.getText(context.Background(), srv.URL+"/x")
	}

	c.mu.RLock()
	tr := c.transport
	c.mu.RUnlock()
	if tr == nil {
		t.Fatal("Transport 应该是复用的，不该为 nil")
	}

	c.Close()

	c.mu.RLock()
	still := c.transport
	c.mu.RUnlock()
	if still != nil {
		t.Error("Close() 后 transport 应被清空")
	}
}

// 改代理配置后 Transport 要重建（否则还是走老代理）
func TestSetProxyRebuildsTransport(t *testing.T) {
	c := NewClient(5*time.Second, "", "off", "auto")

	// 先建一个
	c.getText(context.Background(), "http://127.0.0.1:1/never") // 必然失败，只为触发建 Transport
	c.mu.RLock()
	first := c.transport
	c.mu.RUnlock()
	if first == nil {
		t.Fatal("应该已经建好 Transport")
	}

	// 改代理
	c.SetProxy("http://127.0.0.1:7890", "http")

	c.mu.RLock()
	second := c.transport
	c.mu.RUnlock()
	if second == first {
		t.Error("改代理后 Transport 必须重建，否则还走旧配置")
	}
	if second == nil {
		t.Error("重建后不该是 nil")
	}
}

// 解析进度回调要真的被调用
func TestParseProgressCallback(t *testing.T) {
	const gid, tok, total = "4239999", "abcdef1234", 45

	pages := map[int]string{}
	for p := 0; p < 3; p++ {
		from, to := p*20+1, p*20+20
		if to > total {
			to = total
		}
		page := p
		pages[page] = buildGalleryHTML(gid, from, to, total, page)
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/g/"+gid+"/"+tok+"/", func(w http.ResponseWriter, r *http.Request) {
		p := 0
		if q := r.URL.Query().Get("p"); q != "" {
			fmt.Sscanf(q, "%d", &p)
		}
		h, ok := pages[p]
		if !ok {
			h = buildGalleryHTML(gid, 0, 0, total, 9)
		}
		w.Write([]byte(h))
	})
	srv := httptest.NewServer(mux)
	defer srv.Close()

	old := EHBase
	EHBase = srv.URL
	defer func() { EHBase = old }()

	c := NewClient(5*time.Second, "", "off", "auto")

	var got []int
	var maxDone int
	g, err := c.FetchGalleryPage(context.Background(), gid, tok, func(done, tot, page int) {
		got = append(got, done)
		if done > maxDone {
			maxDone = done
		}
	})
	if err != nil {
		t.Fatalf("解析失败: %v", err)
	}
	if len(got) < 3 {
		t.Errorf("应该至少回调 3 次（翻 3 屏），实际 %d 次", len(got))
	}
	if maxDone != total {
		t.Errorf("最后一次进度应达到 %d，实际 %d", total, maxDone)
	}
	if g.TotalPages != total {
		t.Errorf("总页数不对: %d", g.TotalPages)
	}
	t.Logf("进度回调 %d 次: %v", len(got), got)
}

// 传 nil 回调也不该 panic
func TestParseProgressNil(t *testing.T) {
	const gid, tok, total = "4239998", "abcdef1234", 20

	mux := http.NewServeMux()
	mux.HandleFunc("/g/"+gid+"/"+tok+"/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(buildGalleryHTML(gid, 1, 20, total, 0)))
	})
	srv := httptest.NewServer(mux)
	defer srv.Close()

	old := EHBase
	EHBase = srv.URL
	defer func() { EHBase = old }()

	c := NewClient(5*time.Second, "", "off", "auto")
	if _, err := c.FetchGalleryPage(context.Background(), gid, tok, nil); err != nil {
		t.Errorf("nil 回调不该出错: %v", err)
	}
}