rawdoc/fetch_test.go at main · RandomCodeSpace/rawdoc · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
package main

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

func TestTier1FetchSuccess(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ua := r.Header.Get("User-Agent")
		if ua == "Go-http-client/1.1" || ua == "" {
			t.Errorf("expected browser-like User-Agent, got %q", ua)
		}
		accept := r.Header.Get("Accept")
		if accept == "" {
			t.Error("expected Accept header to be sent")
		}
		fmt.Fprint(w, "<html><body>Hello</body></html>")
	}))
	defer ts.Close()

	opts := &fetchOptions{
		timeout:    5 * time.Second,
		maxRetries: 3,
	}
	result, err := fetchTier1(ts.URL, opts)
	if err != nil {
		t.Fatalf("expected success, got error: %v", err)
	}
	if result.tier != 1 {
		t.Errorf("expected tier 1, got %d", result.tier)
	}
	if result.html == "" {
		t.Error("expected non-empty html")
	}
}

func TestTier1FetchRetryOn5xx(t *testing.T) {
	var attempts int32

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		n := atomic.AddInt32(&attempts, 1)
		if n < 3 {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		fmt.Fprint(w, "<html><body>OK</body></html>")
	}))
	defer ts.Close()

	opts := &fetchOptions{
		timeout:    5 * time.Second,
		maxRetries: 3,
	}
	result, err := fetchTier1(ts.URL, opts)
	if err != nil {
		t.Fatalf("expected success after retries, got error: %v", err)
	}
	if result.tier != 1 {
		t.Errorf("expected tier 1, got %d", result.tier)
	}
	if atomic.LoadInt32(&attempts) != 3 {
		t.Errorf("expected 3 attempts, got %d", atomic.LoadInt32(&attempts))
	}
}

func TestTier1FetchTimeout(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(2 * time.Second)
		fmt.Fprint(w, "<html><body>Late</body></html>")
	}))
	defer ts.Close()

	opts := &fetchOptions{
		timeout:    100 * time.Millisecond,
		maxRetries: 1,
	}
	_, err := fetchTier1(ts.URL, opts)
	if err == nil {
		t.Fatal("expected timeout error, got nil")
	}
}

func TestFetchCustomHeaders(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		val := r.Header.Get("X-Custom")
		if val != "test-value" {
			t.Errorf("expected X-Custom=test-value, got %q", val)
		}
		fmt.Fprint(w, "<html><body>OK</body></html>")
	}))
	defer ts.Close()

	opts := &fetchOptions{
		timeout:    5 * time.Second,
		maxRetries: 3,
		headers:    []string{"X-Custom=test-value"},
	}
	_, err := fetchTier1(ts.URL, opts)
	if err != nil {
		t.Fatalf("expected success, got error: %v", err)
	}
}

func TestTierEscalationOn403(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusForbidden)
	}))
	defer ts.Close()

	opts := &fetchOptions{
		timeout:    5 * time.Second,
		maxRetries: 3,

	}
	_, err := fetch(ts.URL, opts)
	if err == nil {
		t.Fatal("expected error on 403, got nil")
	}
}