Implement unit tests for public packages

This commit adds a comprehensive unit test suite for all public packages.

cmd/clientsim/main.go

@@ -115,13 +115,13 @@ func (c *clientSimulator) Run(ctx context.Context) error {
 	generatorErrors := make(chan error)
 
 	go func() {
-		err := c.framer.ReadFrames(os.Stdin, c.framesIn)
+		err := c.framer.ReadFrames(ctx, os.Stdin, c.framesIn)
 
 		framerErrors <- err
 	}()
 
 	go func() {
-		err := c.framer.WriteFrames(os.Stdout, c.framesOut)
+		err := c.framer.WriteFrames(ctx, os.Stdout, c.framesOut)
 
 		framerErrors <- err
 	}()
@@ -129,7 +129,7 @@ func (c *clientSimulator) Run(ctx context.Context) error {
 	go func() {
 		clientProtocol := protocol.NewClient(c.clientHandler, c.commandGenerator.commands, c.framesIn, c.framesOut, c.logger)
 
-		err := clientProtocol.Handle()
+		err := clientProtocol.Handle(ctx)
 
 		protocolErrors <- err
 	}()
@@ -170,7 +170,7 @@ type ClientHandler struct {
 	logger *logrus.Logger
 }
 
-func (c *ClientHandler) Send(command *protocol.Command, out chan []byte) error {
+func (c *ClientHandler) Send(ctx context.Context, command *protocol.Command, out chan []byte) error {
 	var (
 		frame []byte
 		err   error
@@ -185,7 +185,12 @@ func (c *ClientHandler) Send(command *protocol.Command, out chan []byte) error {
 
 	c.logger.Trace("writing command announcement")
 
-	out <- frame
+	select {
+	case <-ctx.Done():
+		return nil
+	case out <- frame:
+		break
+	}
 
 	frame, err = msgpack.Marshal(command.Command)
 	if err != nil {
@@ -194,17 +199,24 @@ func (c *ClientHandler) Send(command *protocol.Command, out chan []byte) error {
 
 	c.logger.WithField("command", command.Command).Info("write command data")
 
-	out <- frame
+	select {
+	case <-ctx.Done():
+		return nil
+	case out <- frame:
+		break
+	}
 
 	return nil
 }
 
-func (c *ClientHandler) ResponseAnnounce(in chan []byte) (*protocol.Response, error) {
+func (c *ClientHandler) ResponseAnnounce(ctx context.Context, in chan []byte) (*protocol.Response, error) {
 	response := &protocol.Response{}
 
 	var announce messages.ResponseAnnounce
 
 	select {
+	case <-ctx.Done():
+		return nil, nil
 	case frame := <-in:
 		if err := msgpack.Unmarshal(frame, &announce); err != nil {
 			return nil, fmt.Errorf("could not unmarshal response announcement: %w", err)
@@ -220,8 +232,10 @@ func (c *ClientHandler) ResponseAnnounce(in chan []byte) (*protocol.Response, er
 	}
 }
 
-func (c *ClientHandler) ResponseData(in chan []byte, response *protocol.Response) error {
+func (c *ClientHandler) ResponseData(ctx context.Context, in chan []byte, response *protocol.Response) error {
 	select {
+	case <-ctx.Done():
+		return nil
 	case frame := <-in:
 		switch response.Announce.Code {
 		case messages.RespHealth:
@@ -248,7 +262,7 @@ func (c *ClientHandler) ResponseData(in chan []byte, response *protocol.Response
 	return nil
 }
 
-func (c *ClientHandler) HandleResponse(response *protocol.Response) error {
+func (c *ClientHandler) HandleResponse(_ context.Context, response *protocol.Response) error {
 	c.logger.WithField("response", response.Announce).Info("handled response")
 	c.logger.WithField("response", response).Debug("full response")
 

internal/handler/msgpack.go

@@ -18,10 +18,10 @@ limitations under the License.
 package handler
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"math/big"
-	"sync"
 	"time"
 
 	"github.com/shamaton/msgpackgen/msgpack"
@@ -43,35 +43,33 @@ type MsgPackHandler struct {
 	logger          *logrus.Logger
 	healthHandler   *health.Handler
 	fetchCRLHandler *revoking.FetchCRLHandler
-	lock            sync.Mutex
 }
 
-func (m *MsgPackHandler) CommandAnnounce(frames chan []byte) (*protocol.Command, error) {
-	m.lock.Lock()
-	defer m.lock.Unlock()
+func (m *MsgPackHandler) CommandAnnounce(ctx context.Context, frames chan []byte) (*protocol.Command, error) {
+	select {
+	case <-ctx.Done():
+		return nil, nil
+	case frame := <-frames:
+		var ann messages.CommandAnnounce
 
-	frame := <-frames
+		if err := msgpack.Unmarshal(frame, &ann); err != nil {
+			return nil, fmt.Errorf("could not unmarshal command announcement: %w", err)
+		}
 
-	var ann messages.CommandAnnounce
+		if ann.Code == messages.CmdUndef {
+			return nil, fmt.Errorf("received undefined command announcement: %s", ann)
+		}
 
-	if err := msgpack.Unmarshal(frame, &ann); err != nil {
-		return nil, fmt.Errorf("could not unmarshal command announcement: %w", err)
-	}
+		m.logger.WithField("announcement", &ann).Debug("received command announcement")
 
-	if ann.Code == messages.CmdUndef {
-		return nil, fmt.Errorf("received undefined command announcement: %s", ann)
+		return &protocol.Command{Announce: &ann}, nil
 	}
-
-	m.logger.WithField("announcement", &ann).Debug("received command announcement")
-
-	return &protocol.Command{Announce: &ann}, nil
 }
 
-func (m *MsgPackHandler) CommandData(frames chan []byte, command *protocol.Command) error {
-	m.lock.Lock()
-	defer m.lock.Unlock()
-
+func (m *MsgPackHandler) CommandData(ctx context.Context, frames chan []byte, command *protocol.Command) error {
 	select {
+	case <-ctx.Done():
+		return nil
 	case frame := <-frames:
 		err := m.parseCommand(frame, command)
 		if err != nil {
@@ -84,10 +82,7 @@ func (m *MsgPackHandler) CommandData(frames chan []byte, command *protocol.Comma
 	}
 }
 
-func (m *MsgPackHandler) HandleCommand(command *protocol.Command) (*protocol.Response, error) {
-	m.lock.Lock()
-	defer m.lock.Unlock()
-
+func (m *MsgPackHandler) HandleCommand(_ context.Context, command *protocol.Command) (*protocol.Response, error) {
 	var (
 		response *protocol.Response
 		err      error
@@ -110,10 +105,7 @@ func (m *MsgPackHandler) logCommandResponse(command *protocol.Command, response
 	m.logger.WithField("command", command).WithField("response", response).Debug("command and response")
 }
 
-func (m *MsgPackHandler) Respond(response *protocol.Response, out chan []byte) error {
-	m.lock.Lock()
-	defer m.lock.Unlock()
-
+func (m *MsgPackHandler) Respond(ctx context.Context, response *protocol.Response, out chan []byte) error {
 	announce, err := msgpack.Marshal(response.Announce)
 	if err != nil {
 		return fmt.Errorf("could not marshal response announcement: %w", err)
@@ -121,7 +113,12 @@ func (m *MsgPackHandler) Respond(response *protocol.Response, out chan []byte) e
 
 	m.logger.WithField("length", len(announce)).Debug("write response announcement")
 
-	out <- announce
+	select {
+	case <-ctx.Done():
+		return nil
+	case out <- announce:
+		break
+	}
 
 	data, err := msgpack.Marshal(response.Response)
 	if err != nil {
@@ -130,7 +127,12 @@ func (m *MsgPackHandler) Respond(response *protocol.Response, out chan []byte) e
 
 	m.logger.WithField("length", len(data)).Debug("write response")
 
-	out <- data
+	select {
+	case <-ctx.Done():
+		return nil
+	case out <- data:
+		break
+	}
 
 	return nil
 }

internal/serial/seriallink.go

@@ -63,13 +63,13 @@ func (h *Handler) Run(ctx context.Context) error {
 	protocolErrors, framerErrors := make(chan error), make(chan error)
 
 	go func() {
-		err := h.framer.ReadFrames(h.port, h.framesIn)
+		err := h.framer.ReadFrames(ctx, h.port, h.framesIn)
 
 		framerErrors <- err
 	}()
 
 	go func() {
-		err := h.framer.WriteFrames(h.port, h.framesOut)
+		err := h.framer.WriteFrames(ctx, h.port, h.framesOut)
 
 		framerErrors <- err
 	}()
@@ -77,7 +77,7 @@ func (h *Handler) Run(ctx context.Context) error {
 	go func() {
 		serverProtocol := protocol.NewServer(h.serverHandler, h.framesIn, h.framesOut, h.logger)
 
-		err := serverProtocol.Handle()
+		err := serverProtocol.Handle(ctx)
 
 		protocolErrors <- err
 	}()

pkg/protocol/protocol.go

@@ -20,6 +20,7 @@ package protocol
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -54,25 +55,27 @@ func (r *Response) String() string {
 // ServerHandler is responsible for parsing incoming frames and calling commands
 type ServerHandler interface {
 	// CommandAnnounce handles the initial announcement of a command.
-	CommandAnnounce(chan []byte) (*Command, error)
+	CommandAnnounce(context.Context, chan []byte) (*Command, error)
 	// CommandData handles the command data.
-	CommandData(chan []byte, *Command) error
+	CommandData(context.Context, chan []byte, *Command) error
 	// HandleCommand executes the command, generating a response.
-	HandleCommand(*Command) (*Response, error)
+	HandleCommand(context.Context, *Command) (*Response, error)
 	// Respond generates the response for a command.
-	Respond(*Response, chan []byte) error
+	Respond(context.Context, *Response, chan []byte) error
 }
 
 type ClientHandler interface {
-	Send(*Command, chan []byte) error
-	ResponseAnnounce(chan []byte) (*Response, error)
-	ResponseData(chan []byte, *Response) error
-	HandleResponse(*Response) error
+	Send(context.Context, *Command, chan []byte) error
+	ResponseAnnounce(context.Context, chan []byte) (*Response, error)
+	ResponseData(context.Context, chan []byte, *Response) error
+	HandleResponse(context.Context, *Response) error
 }
 
 var (
-	errCommandExpected  = errors.New("command must not be nil")
-	errResponseExpected = errors.New("response must not be nil")
+	errCommandExpected         = errors.New("command must not be nil")
+	errCommandAnnounceExpected = errors.New("command must have an announcement")
+	errCommandDataExpected     = errors.New("command must have data")
+	errResponseExpected        = errors.New("response must not be nil")
 
 	ErrResponseAnnounceTimeoutExpired = errors.New("response announce timeout expired")
 	ErrResponseDataTimeoutExpired     = errors.New("response data timeout expired")
@@ -115,7 +118,7 @@ type ServerProtocol struct {
 	state   protocolState
 }
 
-func (p *ServerProtocol) Handle() error {
+func (p *ServerProtocol) Handle(ctx context.Context) error {
 	var (
 		command  *Command
 		response *Response
@@ -123,67 +126,107 @@ func (p *ServerProtocol) Handle() error {
 	)
 
 	for {
-		p.logger.Debugf("handling protocol state %s", p.state)
+		select {
+		case <-ctx.Done():
+			close(p.out)
 
-		switch p.state {
-		case cmdAnnounce:
-			command, err = p.handler.CommandAnnounce(p.in)
-			if err != nil {
-				p.logger.WithError(err).Error("could not handle command announce")
-
-				break
+			return nil
+		default:
+			p.logger.Debugf("handling protocol state %s", p.state)
+
+			switch p.state {
+			case cmdAnnounce:
+				command = p.commandAnnounce(ctx)
+			case cmdData:
+				err = p.commandData(ctx, command)
+				if err != nil {
+					return err
+				}
+			case handleCommand:
+				response, err = p.handleCommand(ctx, command)
+				if err != nil {
+					return err
+				}
+			case respond:
+				err = p.respond(ctx, response)
+				if err != nil {
+					return err
+				}
+			default:
+				return fmt.Errorf("unknown protocol state %s", p.state)
 			}
+		}
+	}
+}
 
-			p.state = cmdData
-		case cmdData:
-			if command == nil {
-				return errCommandExpected
-			}
+func (p *ServerProtocol) commandAnnounce(ctx context.Context) *Command {
+	command, err := p.handler.CommandAnnounce(ctx, p.in)
+	if err != nil {
+		p.logger.WithError(err).Error("could not handle command announce")
 
-			err = p.handler.CommandData(p.in, command)
-			if err != nil {
-				p.logger.WithError(err).Error("could not handle command data")
+		return nil
+	}
 
-				p.state = cmdAnnounce
+	p.state = cmdData
 
-				break
-			}
+	return command
+}
 
-			p.state = handleCommand
-		case handleCommand:
-			if command == nil {
-				return errCommandExpected
-			}
+func (p *ServerProtocol) commandData(ctx context.Context, command *Command) error {
+	if command == nil || command.Announce == nil {
+		return errCommandAnnounceExpected
+	}
 
-			response, err = p.handler.HandleCommand(command)
-			if err != nil {
-				p.logger.WithError(err).Error("could not handle command")
+	err := p.handler.CommandData(ctx, p.in, command)
+	if err != nil {
+		p.logger.WithError(err).Error("could not handle command data")
 
-				p.state = cmdAnnounce
+		p.state = cmdAnnounce
 
-				break
-			}
+		return nil
+	}
 
-			p.state = respond
-		case respond:
-			if response == nil {
-				return errResponseExpected
-			}
+	p.state = handleCommand
 
-			err = p.handler.Respond(response, p.out)
-			if err != nil {
-				p.logger.WithError(err).Error("could not respond")
+	return nil
+}
 
-				p.state = cmdAnnounce
+func (p *ServerProtocol) handleCommand(ctx context.Context, command *Command) (*Response, error) {
+	if command == nil || command.Announce == nil || command.Command == nil {
+		return nil, errCommandDataExpected
+	}
 
-				break
-			}
+	response, err := p.handler.HandleCommand(ctx, command)
+	if err != nil {
+		p.logger.WithError(err).Error("could not handle command")
 
-			p.state = cmdAnnounce
-		default:
-			return fmt.Errorf("unknown protocol state %s", p.state)
-		}
+		p.state = cmdAnnounce
+
+		return nil, nil
 	}
+
+	p.state = respond
+
+	return response, nil
+}
+
+func (p *ServerProtocol) respond(ctx context.Context, response *Response) error {
+	if response == nil {
+		return errResponseExpected
+	}
+
+	err := p.handler.Respond(ctx, response, p.out)
+	if err != nil {
+		p.logger.WithError(err).Error("could not respond")
+
+		p.state = cmdAnnounce
+
+		return nil
+	}
+
+	p.state = cmdAnnounce
+
+	return nil
 }
 
 func NewServer(handler ServerHandler, in, out chan []byte, logger *logrus.Logger) *ServerProtocol {
@@ -196,14 +239,6 @@ func NewServer(handler ServerHandler, in, out chan []byte, logger *logrus.Logger
 	}
 }
 
-// Framer handles bytes on the wire by adding or removing framing information.
-type Framer interface {
-	// ReadFrames reads data frames and publishes unframed data to the channel.
-	ReadFrames(io.Reader, chan []byte) error
-	// WriteFrames takes data from the channel and writes framed data to the writer.
-	WriteFrames(io.Writer, chan []byte) error
-}
-
 type ClientProtocol struct {
 	handler  ClientHandler
 	commands chan *Command
@@ -212,80 +247,121 @@ type ClientProtocol struct {
 	state    protocolState
 }
 
-func (p *ClientProtocol) Handle() error {
+func (p *ClientProtocol) Handle(ctx context.Context) error {
 	var (
-		command  *Command
 		response *Response
 		err      error
 	)
 
 	for {
-		p.logger.Debugf("handling protocol state %s", p.state)
+		select {
+		case <-ctx.Done():
+			return nil
+		default:
+			p.logger.Debugf("handling protocol state %s", p.state)
 
-		switch p.state {
-		case cmdAnnounce:
-			command = <-p.commands
-			if command == nil {
-				return errCommandExpected
+			switch p.state {
+			case cmdAnnounce:
+				err = p.cmdAnnounce(ctx)
+				if err != nil {
+					return err
+				}
+			case respAnnounce:
+				response = p.respAnnounce(ctx)
+			case respData:
+				err = p.respData(ctx, response)
+				if err != nil {
+					return err
+				}
+			case handleResponse:
+				err = p.handleResponse(ctx, response)
+				if err != nil {
+					return err
+				}
+			default:
+				return fmt.Errorf("unknown protocol state %s", p.state)
 			}
+		}
+	}
+}
 
-			err = p.handler.Send(command, p.out)
-			if err != nil {
-				p.logger.WithError(err).Error("could not send command announce")
-
-				break
-			}
+func (p *ClientProtocol) cmdAnnounce(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return nil
+	case command := <-p.commands:
+		if command == nil {
+			return errCommandExpected
+		}
 
-			p.state = respAnnounce
-		case respAnnounce:
-			response, err = p.handler.ResponseAnnounce(p.in)
-			if err != nil {
-				p.logger.WithError(err).Error("could not handle response announce")
+		err := p.handler.Send(ctx, command, p.out)
+		if err != nil {
+			p.logger.WithError(err).Error("could not send command announce")
 
-				p.state = cmdAnnounce
+			return nil
+		}
+	}
 
-				break
-			}
+	p.state = respAnnounce
 
-			p.state = respData
-		case respData:
-			if response == nil {
-				return errResponseExpected
-			}
+	return nil
+}
 
-			err = p.handler.ResponseData(p.in, response)
-			if err != nil {
-				p.logger.WithError(err).Error("could not handle response data")
+func (p *ClientProtocol) respAnnounce(ctx context.Context) *Response {
+	response, err := p.handler.ResponseAnnounce(ctx, p.in)
+	if err != nil {
+		p.logger.WithError(err).Error("could not handle response announce")
 
-				if errors.Is(err, ErrResponseDataTimeoutExpired) {
-					p.state = cmdAnnounce
-				} else {
-					p.state = respAnnounce
-				}
+		p.state = cmdAnnounce
 
-				break
-			}
+		return nil
+	}
 
-			p.state = handleResponse
-		case handleResponse:
-			if response == nil {
-				return errResponseExpected
-			}
+	p.state = respData
 
-			err = p.handler.HandleResponse(response)
-			if err != nil {
-				p.logger.WithError(err).Error("could not handle response")
+	return response
+}
 
-				p.state = cmdAnnounce
+func (p *ClientProtocol) respData(ctx context.Context, response *Response) error {
+	if response == nil || response.Announce == nil {
+		return errResponseExpected
+	}
 
-				break
-			}
+	err := p.handler.ResponseData(ctx, p.in, response)
+	if err != nil {
+		p.logger.WithError(err).Error("could not handle response data")
 
+		if errors.Is(err, ErrResponseDataTimeoutExpired) {
 			p.state = cmdAnnounce
-		default:
-			return fmt.Errorf("unknown protocol state %s", p.state)
+		} else {
+			p.state = respAnnounce
 		}
+
+		return nil
+	}
+
+	p.state = handleResponse
+
+	return nil
+}
+
+func (p *ClientProtocol) handleResponse(ctx context.Context, response *Response) error {
+	if response == nil || response.Announce == nil || response.Response == nil {
+		return errResponseExpected
+	}
+
+	err := p.handler.HandleResponse(ctx, response)
+	if err != nil {
+		p.logger.WithError(err).Error("could not handle response")
+
+		p.state = cmdAnnounce
+
+		return nil
 	}
+
+	p.state = cmdAnnounce
+
+	return nil
 }
 
 func NewClient(
@@ -304,6 +380,14 @@ func NewClient(
 	}
 }
 
+// Framer handles bytes on the wire by adding or removing framing information.
+type Framer interface {
+	// ReadFrames reads data frames and publishes unframed data to the channel.
+	ReadFrames(context.Context, io.Reader, chan []byte) error
+	// WriteFrames takes data from the channel and writes framed data to the writer.
+	WriteFrames(context.Context, io.Writer, chan []byte) error
+}
+
 const bufferSize = 1024
 const readInterval = 50 * time.Millisecond
 
@@ -319,7 +403,7 @@ func NewCOBSFramer(logger *logrus.Logger) *COBSFramer {
 	}
 }
 
-func (c *COBSFramer) ReadFrames(reader io.Reader, frameChan chan []byte) error {
+func (c *COBSFramer) ReadFrames(ctx context.Context, reader io.Reader, frameChan chan []byte) error {
 	var (
 		err              error
 		raw, data, frame []byte
@@ -328,49 +412,55 @@ func (c *COBSFramer) ReadFrames(reader io.Reader, frameChan chan []byte) error {
 	buffer := &bytes.Buffer{}
 
 	for {
-		raw, err = c.readRaw(reader)
-		if err != nil {
-			close(frameChan)
+		select {
+		case <-ctx.Done():
+			return nil
+		default:
+			raw, err = c.readRaw(reader)
+			if err != nil {
+				close(frameChan)
 
-			return err
-		}
+				return err
+			}
 
-		if len(raw) == 0 {
-			time.Sleep(readInterval)
+			if len(raw) == 0 {
+				time.Sleep(readInterval)
 
-			continue
-		}
+				continue
+			}
 
-		c.logger.Tracef("read %d raw bytes", len(raw))
+			c.logger.Tracef("read %d raw bytes", len(raw))
 
-		buffer.Write(raw)
+			buffer.Write(raw)
 
-		for {
-			data, err = buffer.ReadBytes(c.config.SpecialByte)
-			if err != nil {
-				if errors.Is(err, io.EOF) {
-					buffer.Write(data)
+			for {
+				data, err = buffer.ReadBytes(c.config.SpecialByte)
+				if err != nil {
+					if errors.Is(err, io.EOF) {
+						buffer.Write(data)
+
+						break
+					}
 
-					break
+					// this is a safety measure, buffer.ReadBytes should only return io.EOF
+					return fmt.Errorf("could not read from buffer: %w", err)
 				}
 
-				return fmt.Errorf("could not read from buffer: %w", err)
-			}
+				if err = cobs.Verify(data, c.config); err != nil {
+					c.logger.WithError(err).Warnf("skipping invalid frame of %d bytes", len(data))
 
-			if err = cobs.Verify(data, c.config); err != nil {
-				c.logger.WithError(err).Warnf("skipping invalid frame of %d bytes", len(data))
+					continue
+				}
 
-				continue
-			}
+				frame = cobs.Decode(data, c.config)
 
-			frame = cobs.Decode(data, c.config)
+				c.logger.Tracef("frame decoded to length %d", len(frame))
 
-			c.logger.Tracef("frame decoded to length %d", len(frame))
+				frameChan <- frame
+			}
 
-			frameChan <- frame
+			c.logger.Tracef("read buffer is now %d bytes long", buffer.Len())
 		}
-
-		c.logger.Tracef("read buffer is now %d bytes long", buffer.Len())
 	}
 }
 
@@ -380,7 +470,7 @@ func (c *COBSFramer) readRaw(reader io.Reader) ([]byte, error) {
 	count, err := reader.Read(buf)
 	if err != nil {
 		if errors.Is(err, io.EOF) {
-			return []byte{}, nil
+			return buf[:count], nil
 		}
 
 		return nil, fmt.Errorf("could not read data: %w", err)
@@ -391,23 +481,26 @@ func (c *COBSFramer) readRaw(reader io.Reader) ([]byte, error) {
 	return raw, nil
 }
 
-func (c *COBSFramer) WriteFrames(writer io.Writer, frameChan chan []byte) error {
+func (c *COBSFramer) WriteFrames(ctx context.Context, writer io.Writer, frameChan chan []byte) error {
 	for {
-		frame := <-frameChan
+		select {
+		case <-ctx.Done():
+			return nil
+		case frame := <-frameChan:
+			if frame == nil {
+				c.logger.Debug("channel closed")
 
-		if frame == nil {
-			c.logger.Debug("channel closed")
+				return nil
+			}
 
-			return nil
-		}
+			encoded := cobs.Encode(frame, c.config)
 
-		encoded := cobs.Encode(frame, c.config)
+			n, err := io.Copy(writer, bytes.NewReader(encoded))
+			if err != nil {
+				return fmt.Errorf("cold not write data: %w", err)
+			}
 
-		n, err := io.Copy(writer, bytes.NewReader(encoded))
-		if err != nil {
-			return fmt.Errorf("cold not write data: %w", err)
+			c.logger.Tracef("wrote %d bytes", n)
 		}
-
-		c.logger.Tracef("wrote %d bytes", n)
 	}
 }