Go语言Worker方法的正确测试方案咨询

先来说你遇到的TestWorkerCloseContext阻塞问题——根源出在你的MockQueue.TakeMessage()实现上：你创建了一个无缓冲通道后直接往里面发数据，无缓冲通道的发送操作会一直卡着，直到有接收方把数据读走。但你的Worker还没走到读取通道的那一步，就先卡在了TakeMessage()调用上，根本没进入后面的select循环，所以你调用cancel()后，Worker根本没机会响应ctx.Done()，测试自然就阻塞了。

第一步：修复MockQueue的实现（正确用testify/mock） #

既然你已经引入了testify/mock，就别硬编码逻辑了，按照mock库的规则来写，既能避免阻塞，又能灵活控制返回值：

func (m *MockQueue) TakeMessage() (<-chan string, error) {
	args := m.Called()
	return args.Get(0).(<-chan string), args.Error(1)
}

第二步：修复TestWorkerCloseContext测试 #

现在我们可以让MockQueue返回一个不会阻塞的空缓冲通道，这样Worker能顺利进入select循环，响应ctx的取消信号：

func TestWorkerCloseContext(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// 初始化mock实例
	mockQueue := new(MockQueue)
	mockDownload := new(MockDownloader)

	// 设置MockQueue返回空的缓冲通道，避免阻塞
	msgChan := make(chan string, 1)
	mockQueue.On("TakeMessage").Return(msgChan, nil)

	worker := NewWorker(mockQueue, mockDownload)
	resultCh := make(chan error, 1) // 用缓冲通道防止goroutine被阻塞

	go func() {
		resultCh <- worker.Worker(ctx)
	}()

	// 触发ctx取消
	cancel()
	// 等待Worker返回结果
	err := <-resultCh
	assert.Nil(t, err)

	// 验证mock方法是否按预期被调用
	mockQueue.AssertExpectations(t)
}

接下来是消息传递逻辑的测试用例实现，我们要验证当Queue返回消息时，Worker会正确调用Download方法处理这个消息，用mock库的断言就能轻松做到：

func TestWorkerMessageReceive(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	testURL := "some_url/some.txt"

	// 初始化mock
	mockQueue := new(MockQueue)
	mockDownload := new(MockDownloader)

	// 给MockQueue设置返回包含测试URL的缓冲通道
	msgChan := make(chan string, 1)
	msgChan <- testURL
	mockQueue.On("TakeMessage").Return(msgChan, nil)

	// 预先声明期望：Download方法会被调用一次，参数是testURL
	mockDownload.On("Download", testURL).Return(nil)

	worker := NewWorker(mockQueue, mockDownload)
	resultCh := make(chan error, 1)

	go func() {
		err := worker.Worker(ctx)
		resultCh <- err
	}()

	// 确认消息被Worker读取（也可以用sync.WaitGroup更严谨）
	<-msgChan
	// 触发Worker退出
	cancel()

	// 验证结果
	err := <-resultCh
	assert.Nil(t, err)
	// 检查Download是否按预期被调用
	mockDownload.AssertExpectations(t)
	mockQueue.AssertExpectations(t)
}

额外优化小技巧 #

• 测试里的通道尽量用缓冲通道，避免不必要的阻塞；
• 如果要更精准地控制并发流程（比如等Download执行完再取消ctx），可以用sync.WaitGroup：

  // 在TestWorkerMessageReceive中修改mock期望：
  var wg sync.WaitGroup
  wg.Add(1)
  mockDownload.On("Download", testURL).Return(nil).Run(func(args mock.Arguments) {
      wg.Done() // 当Download被调用时标记任务完成
  })
  
  // 启动Worker后，等待Download执行完毕再取消ctx
  wg.Wait()
  cancel()
  

内容的提问来源于stack exchange，提问作者vsfm80