编程语言

Go教程(五)实战 —— CLI 工具与测试

Go 05:实战 —— CLI 工具与测试

Go 写 CLI 工具是它最舒服的场景之一:启动快、单文件分发、跨平台。很多知名工具都是 Go 写的:docker、kubectl、gh、terraform、hugo。

本篇做一个 URL 健康检查工具 urlcheck,综合运用前面学到的内容,然后为它编写完整测试。

1. 需求

urlcheck -f urls.txt -c 10 -timeout 3s -o report.json
  • 从文件读取 URL 列表(每行一个,# 开头为注释)
  • 并发检查(并发数可配置),记录状态码与耗时
  • 终端输出表格,可选输出 JSON 报告
  • 有失败时退出码为 1(方便接入 CI / crontab 告警)

2. 项目结构

urlcheck/
├── go.mod
├── main.go               # 参数解析与输出
└── checker/
    ├── checker.go        # 核心逻辑(可测试)
    └── checker_test.go

原则:main 只做"胶水",逻辑放到独立包里,便于测试。

3. 核心逻辑

checker/checker.go:

package checker

import (
	"bufio"
	"context"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"
)

type Result struct {
	URL      string        `json:"url"`
	Status   int           `json:"status"`
	Duration time.Duration `json:"durationNs"`
	Error    string        `json:"error,omitempty"`
}

func (r Result) OK() bool {
	return r.Error == "" && r.Status >= 200 && r.Status < 400
}

// ParseURLs 从任意 io.Reader 读取 URL(文件、stdin、字符串都可以)
func ParseURLs(r io.Reader) ([]string, error) {
	var urls []string
	sc := bufio.NewScanner(r)
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		urls = append(urls, line)
	}
	return urls, sc.Err()
}

type Checker struct {
	Client      *http.Client
	Concurrency int
}

func (c *Checker) checkOne(ctx context.Context, url string) Result {
	start := time.Now()
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return Result{URL: url, Error: err.Error()}
	}
	resp, err := c.Client.Do(req)
	if err != nil {
		return Result{URL: url, Error: err.Error(), Duration: time.Since(start)}
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, resp.Body)   // 读完 body 以便连接复用
	return Result{URL: url, Status: resp.StatusCode, Duration: time.Since(start)}
}

// CheckAll 并发检查,结果顺序与输入一致
func (c *Checker) CheckAll(ctx context.Context, urls []string) []Result {
	results := make([]Result, len(urls))
	sem := make(chan struct{}, max(1, c.Concurrency))   // 信号量控制并发
	var wg sync.WaitGroup

	for i, u := range urls {
		wg.Add(1)
		go func() {
			defer wg.Done()
			sem <- struct{}{}
			defer func() { <-sem }()
			results[i] = c.checkOne(ctx, u)   // 每个 goroutine 写不同下标,无需加锁
		}()
	}
	wg.Wait()
	return results
}

4. 命令行入口

main.go:

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"text/tabwriter"
	"time"

	"example.com/urlcheck/checker"
)

func main() {
	file := flag.String("f", "", "file containing urls (default: stdin)")
	concurrency := flag.Int("c", 5, "concurrency")
	timeout := flag.Duration("timeout", 5*time.Second, "timeout per request")
	output := flag.String("o", "", "write json report to file")
	flag.Parse()

	in := os.Stdin
	if *file != "" {
		f, err := os.Open(*file)
		if err != nil {
			fmt.Fprintln(os.Stderr, "open:", err)
			os.Exit(2)
		}
		defer f.Close()
		in = f
	}

	urls, err := checker.ParseURLs(in)
	if err != nil {
		fmt.Fprintln(os.Stderr, "read:", err)
		os.Exit(2)
	}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)   // Ctrl+C 可中断
	defer stop()

	c := &checker.Checker{Client: &http.Client{Timeout: *timeout}, Concurrency: *concurrency}
	results := c.CheckAll(ctx, urls)

	tw := tabwriter.NewWriter(os.Stdout, 0, 4, 2, ' ', 0)
	fmt.Fprintln(tw, "RESULT\tSTATUS\tTIME\tURL")
	failed := 0
	for _, r := range results {
		mark := "OK"
		if !r.OK() {
			mark, failed = "FAIL", failed+1
		}
		fmt.Fprintf(tw, "%s\t%d\t%v\t%s\n", mark, r.Status, r.Duration.Round(time.Millisecond), r.URL)
	}
	tw.Flush()
	fmt.Printf("\ntotal=%d failed=%d\n", len(results), failed)

	if *output != "" {
		data, _ := json.MarshalIndent(results, "", "  ")
		if err := os.WriteFile(*output, data, 0o644); err != nil {
			fmt.Fprintln(os.Stderr, "write report:", err)
		}
	}
	if failed > 0 {
		os.Exit(1)
	}
}
go mod init example.com/urlcheck
go run . -f urls.txt -c 10
cat urls.txt | go run .          # 也支持管道输入

子命令较多的复杂 CLI(如 tool init、tool run)推荐 spf13/cobra(≈ Python 的 click/typer)。

5. 测试

Go 内置测试框架,无需安装:文件以 _test.go 结尾,函数以 Test 开头。

pytest go test
def test_xxx(): func TestXxx(t *testing.T)
assert a == b if a != b { t.Errorf(...) }(或用 testify/assert)
@pytest.mark.parametrize 表驱动测试 + t.Run 子测试
fixture 辅助函数 +t.Cleanup
responses / mock server net/http/httptest
pytest-benchmark func BenchmarkXxx(b *testing.B)
pytest --cov go test -cover

checker/checker_test.go:

package checker

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

// 表驱动测试 ≈ pytest.mark.parametrize
func TestParseURLs(t *testing.T) {
	tests := []struct {
		name  string
		input string
		want  int
	}{
		{"empty", "", 0},
		{"single", "https://a.com", 1},
		{"skip comment and blank", "# comment\n\nhttps://a.com\n  https://b.com  ", 2},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := ParseURLs(strings.NewReader(tt.input))
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}
			if len(got) != tt.want {
				t.Errorf("got %d urls, want %d: %v", len(got), tt.want, got)
			}
		})
	}
}

// 用 httptest 启动本地假服务器,不依赖外网
func newTestServer(t *testing.T) *httptest.Server {
	t.Helper()
	mux := http.NewServeMux()
	mux.HandleFunc("/ok", func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(200) })
	mux.HandleFunc("/error", func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(500) })
	mux.HandleFunc("/slow", func(w http.ResponseWriter, _ *http.Request) {
		time.Sleep(200 * time.Millisecond)
	})
	srv := httptest.NewServer(mux)
	t.Cleanup(srv.Close)        // ≈ fixture 的 yield 之后的清理
	return srv
}

func TestCheckAll(t *testing.T) {
	srv := newTestServer(t)
	c := &Checker{Client: &http.Client{Timeout: 100 * time.Millisecond}, Concurrency: 2}

	urls := []string{srv.URL + "/ok", srv.URL + "/error", srv.URL + "/slow", "http://invalid.localhost:1"}
	results := c.CheckAll(context.Background(), urls)

	want := []bool{true, false, false, false}
	for i, r := range results {
		if r.URL != urls[i] {
			t.Errorf("result order mismatch at %d", i)
		}
		if r.OK() != want[i] {
			t.Errorf("%s: OK()=%v, want %v (status=%d err=%s)", r.URL, r.OK(), want[i], r.Status, r.Error)
		}
	}
}

func TestCheckAllCancel(t *testing.T) {
	srv := newTestServer(t)
	c := &Checker{Client: http.DefaultClient, Concurrency: 1}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()   // 提前取消

	for _, r := range c.CheckAll(ctx, []string{srv.URL + "/slow"}) {
		if r.Error == "" {
			t.Error("expected error after cancel")
		}
	}
}

// 基准测试
func BenchmarkParseURLs(b *testing.B) {
	input := strings.Repeat("https://example.com/path\n", 1000)
	for b.Loop() {          // Go 1.24+;旧版本用 for i := 0; i < b.N; i++
		_, _ = ParseURLs(strings.NewReader(input))
	}
}

运行:

go test ./...                      # 全部测试
go test -v -run TestParseURLs ./checker   # 指定测试 ≈ pytest -k
go test -race ./...                # 数据竞争检测
go test -cover ./...               # 覆盖率
go test -bench . -benchmem ./checker      # 基准测试

测试 HTTP Handler(对应 04 篇的 API)

func TestCreateRecord(t *testing.T) {
	mux := http.NewServeMux()
	handler.NewRecordHandler(store.NewMemoryStore()).Register(mux)

	body := strings.NewReader(`{"type":"expense","category":"food","amount":10,"date":"2026-11-01"}`)
	req := httptest.NewRequest(http.MethodPost, "/api/records", body)
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)

	if rec.Code != http.StatusCreated {
		t.Fatalf("status = %d, body = %s", rec.Code, rec.Body)
	}
}

想要更像 pytest 的断言风格,可用 github.com/stretchr/testify/assert:assert.Equal(t, 201, rec.Code)。

6. 发布:交叉编译 + Docker

# 一台 Mac 上编译出所有平台的二进制
GOOS=linux   GOARCH=amd64 go build -ldflags="-s -w" -o dist/urlcheck-linux .
GOOS=windows GOARCH=amd64 go build -o dist/urlcheck.exe .
GOOS=darwin  GOARCH=arm64 go build -o dist/urlcheck-mac .

多阶段 Dockerfile(最终镜像只有十几 MB):

FROM golang:1.24 AS build
WORKDIR /src
COPY go.mod go.sum* ./
RUN go mod download
COPY . .
RUN CGO_ENABLED=0 go build -ldflags="-s -w" -o /urlcheck .

FROM gcr.io/distroless/static
COPY --from=build /urlcheck /urlcheck
ENTRYPOINT ["/urlcheck"]

7. 更多适合用 Go 做的实际场景

场景 说明 相关库
运维小工具 日志分析、批量 SSH、文件同步 golang.org/x/crypto/ssh
测试辅助服务 Mock Server、压测工具、数据造数 net/http、vegeta
API 网关 / 反向代理 高并发转发 net/http/httputil.ReverseProxy
消息消费者 消费 Kafka/RabbitMQ 并处理 segmentio/kafka-go
定时任务服务 替代 crontab + Python 脚本 robfig/cron
K8s Operator / 云原生 扩展 Kubernetes controller-runtime

8. 练习

  1. 给 urlcheck 增加 -retry 参数:失败时重试 N 次,并为重试逻辑编写测试。
  2. 增加 -expect 参数,校验响应体包含指定字符串。
  3. 用 cobra 改造成两个子命令:urlcheck run 和 urlcheck serve(以 HTTP 接口形式提供检查服务)。
  4. 写一个简单压测工具:-n 1000 -c 50 URL,输出 QPS、P50/P95/P99 延迟。

文章评论 (0)