naza/pkg/mockwriter/mockwriter.go

package mockwriter

import (
	"bytes"
	"errors"
)

// TODO chef: 可以添加一个接口，获取内部 buffer 的内容

type WriterType uint8

const (
	WriterTypeDoNothing WriterType = iota
	WriterTypeReturnError
	WriterTypeIntoBuffer
)

var (
	mockWriterErr = errors.New("mockwriter: a mock error")
)

type MockWriter struct {
	t     WriterType
	ts    map[uint32]WriterType
	count uint32
	b bytes.Buffer
}

func NewMockWriter(t WriterType) *MockWriter {
	return &MockWriter{
		t: t,
	}
}

// 为某些写操作指定特定的类型，次数从 0 开始计数
func (w *MockWriter) SetSpecificType(ts map[uint32]WriterType) {
	w.ts = ts
}

func (w *MockWriter) Write(b []byte) (int, error) {
	t, exist := w.ts[w.count]
	w.count++
	if !exist {
		t = w.t
	}
	switch t {
	case WriterTypeDoNothing:
		return len(b), nil
	case WriterTypeReturnError:
		return 0, mockWriterErr
	//case WriterTypeIntoBuffer:
	//	return w.b.Write(b)
	}

	return w.b.Write(b)

	//panic("never reach here.")
}