hash-of-wisdom/internal/protocol/codec_test.go

package protocol

import (
	"bytes"
	"io"
	"testing"
	"time"

	"hash-of-wisdom/internal/pow/challenge"
	"hash-of-wisdom/internal/quotes"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestCodec_Decode(t *testing.T) {
	codec := NewCodec()

	tests := []struct {
		name    string
		data    []byte
		wantErr string
	}{
		{
			name:    "empty data",
			data:    []byte{},
			wantErr: "EOF",
		},
		{
			name:    "invalid message type",
			data:    []byte{0xFF, 0x00, 0x00, 0x00, 0x00},
			wantErr: "invalid message type",
		},
		{
			name:    "incomplete header",
			data:    []byte{0x01, 0x00, 0x00},
			wantErr: "failed to read payload length",
		},
		{
			name:    "payload too large",
			data:    append([]byte{0x01}, encodeBigEndianUint32(MaxPayloadSize+1)...),
			wantErr: "payload length",
		},
		{
			name:    "incomplete payload",
			data:    []byte{0x01, 0x00, 0x00, 0x00, 0x05, 0x01, 0x02},
			wantErr: "failed to read payload",
		},
		{
			name:    "invalid UTF-8 in payload",
			data:    []byte{0x01, 0x00, 0x00, 0x00, 0x03, 0xFF, 0xFE, 0xFD},
			wantErr: "invalid UTF-8",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			buf := bytes.NewBuffer(tt.data)
			_, err := codec.Decode(buf)
			assert.Error(t, err)
			assert.Contains(t, err.Error(), tt.wantErr)
		})
	}
}

func TestCodec_ReadError_Handling(t *testing.T) {
	codec := NewCodec()

	// Create a reader that fails after reading header
	reader := &failingReader{
		data:      []byte{0x01, 0x00, 0x00, 0x00, 0x05},
		failAfter: 5,
	}

	_, err := codec.Decode(reader)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "failed to read payload")
}

func TestChallengeResponse_Encode(t *testing.T) {
	challenge := &challenge.Challenge{
		Timestamp:  time.Now().Unix(),
		Difficulty: 4,
		Resource:   "quotes",
		Random:     []byte("random123"),
		HMAC:       []byte("hmac_signature"),
	}

	response := &ChallengeResponse{Challenge: challenge}

	var buf bytes.Buffer
	err := response.Encode(&buf)
	require.NoError(t, err)

	// Verify the encoded message can be decoded
	codec := NewCodec()
	decoded, err := codec.Decode(&buf)
	require.NoError(t, err)

	assert.Equal(t, ChallengeResponseType, decoded.Type)
	assert.Contains(t, string(decoded.Payload), "quotes")
	assert.Contains(t, string(decoded.Payload), "cmFuZG9tMTIz") // "random123" base64 encoded
}

func TestSolutionResponse_Encode(t *testing.T) {
	quote := &quotes.Quote{
		Text:   "Test quote",
		Author: "Test author",
	}

	response := &SolutionResponse{Quote: quote}

	var buf bytes.Buffer
	err := response.Encode(&buf)
	require.NoError(t, err)

	// Verify the encoded message can be decoded
	codec := NewCodec()
	decoded, err := codec.Decode(&buf)
	require.NoError(t, err)

	assert.Equal(t, QuoteResponseType, decoded.Type)
	assert.Contains(t, string(decoded.Payload), "Test quote")
	assert.Contains(t, string(decoded.Payload), "Test author")
}

func TestErrorResponse_Encode(t *testing.T) {
	errorResp := &ErrorResponse{
		Code:       "INVALID_SOLUTION",
		Message:    "The provided PoW solution is incorrect",
		RetryAfter: 30,
		Details:    map[string]string{"attempt": "1"},
	}

	var buf bytes.Buffer
	err := errorResp.Encode(&buf)
	require.NoError(t, err)

	// Verify the encoded message can be decoded
	codec := NewCodec()
	decoded, err := codec.Decode(&buf)
	require.NoError(t, err)

	assert.Equal(t, ErrorResponseType, decoded.Type)
	assert.Contains(t, string(decoded.Payload), "INVALID_SOLUTION")
	assert.Contains(t, string(decoded.Payload), "The provided PoW solution is incorrect")
	assert.Contains(t, string(decoded.Payload), "30")
}

func TestResponse_WriteError_Handling(t *testing.T) {
	response := &ErrorResponse{
		Code:    "TEST_ERROR",
		Message: "Test message",
	}

	// Create a writer that fails immediately
	writer := &failingWriter{failAfter: 1}

	err := response.Encode(writer)
	assert.Error(t, err)
}

// Helper functions and types for testing

func encodeBigEndianUint32(val uint32) []byte {
	return []byte{
		byte(val >> 24),
		byte(val >> 16),
		byte(val >> 8),
		byte(val),
	}
}

type failingWriter struct {
	written   int
	failAfter int
}

func (w *failingWriter) Write(data []byte) (int, error) {
	if w.written >= w.failAfter {
		return 0, io.ErrShortWrite
	}

	remaining := w.failAfter - w.written
	if len(data) <= remaining {
		w.written += len(data)
		return len(data), nil
	}

	w.written = w.failAfter
	return remaining, io.ErrShortWrite
}

type failingReader struct {
	data      []byte
	pos       int
	failAfter int
}

func (r *failingReader) Read(buf []byte) (int, error) {
	if r.pos >= r.failAfter {
		return 0, io.ErrUnexpectedEOF
	}

	if r.pos >= len(r.data) {
		return 0, io.EOF
	}

	n := copy(buf, r.data[r.pos:])
	r.pos += n

	if r.pos >= r.failAfter {
		return n, io.ErrUnexpectedEOF
	}

	return n, nil
}