Stackoverflow: 79205504

main.go

package main

import (
	"fmt"

	myproto "github.com/DazWilkin/stackoverflow/79205504/protos"
	"google.golang.org/protobuf/proto"
)

type Messager interface {
	Header() byte
	Encode() ([]byte, error)
}

var (
	_ Messager = (*EmptyMessage)(nil)
	_ Messager = (*Message1)(nil)
	_ Messager = (*Message2)(nil)
)

type EmptyMessage struct{}

func (m *EmptyMessage) Header() byte {
	return 0
}
func (m *EmptyMessage) Encode() ([]byte, error) {
	return []byte{m.Header()}, nil
}

type Message1 struct {
	body *myproto.FirstMessage
}

func (m *Message1) Header() byte {
	return 1
}
func (m *Message1) Encode() ([]byte, error) {
	b, err := proto.Marshal(m.body)
	if err != nil {
		return nil, err
	}

	return append([]byte{m.Header()}, b...), nil
}

type Message2 struct {
	body *myproto.SecondMessage
}

func (m *Message2) Header() byte {
	return 2
}
func (m *Message2) Encode() ([]byte, error) {
	b, err := proto.Marshal(m.body)
	if err != nil {
		return nil, err
	}

	return append([]byte{m.Header()}, b...), nil
}

func NewMessage(b []byte) (Messager, error) {
	if len(b) == 0 {
		return nil, fmt.Errorf("unable to decode message")
	}
	switch b[0] {
	case 0:
		return &EmptyMessage{}, nil
	case 1:
		body := &myproto.FirstMessage{}
		err := proto.Unmarshal(b[1:], body)
		if err != nil {
			return nil, err
		}

		return &Message1{
			body: body,
		}, nil
	case 2:
		body := &myproto.SecondMessage{}
		err := proto.Unmarshal(b[1:], body)
		if err != nil {
			return nil, err
		}

		return &Message2{
			body: body,
		}, nil
	default:
		return nil, fmt.Errorf("unable to decode message")
	}
}

main_test.go

package main

import (
	"fmt"
	"reflect"
	"testing"

	myproto "github.com/DazWilkin/stackoverflow/79205504/protos"
	"google.golang.org/protobuf/proto"
)

// TestNewMessage tests the NewMessage function
// Message[1|2] have body: proto.Message and so can't be reflect.DeepEqual'd
// Instead since there's no other fields, using proto.Equal on the body's
func TestNewMessage(t *testing.T) {
	for i, test := range []struct {
		data []byte
		want Messager
	}{
		{
			data: []byte{0},
			want: &EmptyMessage{},
		},
		{
			data: append([]byte{1, 10, 7}, []byte("Freddie")...),
			want: &Message1{
				body: &myproto.FirstMessage{
					Dog: "Freddie",
				},
			},
		},
		{
			data: append([]byte{2, 10, 6}, []byte("Emmett")...),
			want: &Message2{
				body: &myproto.SecondMessage{
					Cat: "Emmett",
				},
			},
		},
	} {
		t.Run(
			fmt.Sprintf("test: %d", i),
			func(t *testing.T) {
				fmt.Printf("data: %x\n", test.data)
				m, err := NewMessage(test.data)
				if err != nil {
					t.Error(err)
				}
				switch got := m.(type) {
				case *EmptyMessage:
					want := test.want.(*EmptyMessage)
					if !reflect.DeepEqual(got, want) {
						t.Errorf("got:\n%v\nwant:\n%v", got, want)
					}
				case *Message1:
					want := test.want.(*Message1)
					if !proto.Equal(got.body, want.body) {
						t.Errorf("got:\n%v\nwant:\n%v", got, want)
					}
				case *Message2:
					want := test.want.(*Message2)
					if !proto.Equal(got.body, want.body) {
						t.Errorf("got:\n%v\nwant:\n%v", got, want)
					}
				}
			},
		)
	}
}

// TestEncode tests the encode methods
func TestEncoded(t *testing.T) {
	for i, test := range []struct {
		data Messager
		want []byte
	}{
		{
			data: &EmptyMessage{},
			want: []byte{0},
		},
		{
			data: &Message1{
				body: &myproto.FirstMessage{
					Dog: "Freddie",
				},
			},
			want: append([]byte{1, 10, 7}, []byte("Freddie")...),
		},
		{
			data: &Message2{
				body: &myproto.SecondMessage{
					Cat: "Emmett",
				},
			},
			want: append([]byte{2, 10, 6}, []byte("Emmett")...),
		},
	} {
		t.Run(
			fmt.Sprintf("test: %d", i),
			func(t *testing.T) {
				fmt.Printf("data: %x\n", test.data)
				b, err := test.data.Encode()
				if err != nil {
					t.Error(err)
				}
				if !reflect.DeepEqual(b, test.want) {
					t.Errorf("got:\n%x\nwant:\n%x", b, test.want)
				}
			},
		)
	}
}

my.proto

syntax = "proto3";

// FirstMessage corresponds to MessageType::Message1 (!)
message FirstMessage {
    string dog = 1;
}
// SecondMessage correponds to MessageType::Message2 (!)
message SecondMessage {
    string cat = 1;
}
// ThirdMessage corresponds to no MessageType::{Invalid}
message ThirdMessage {
    string rabbit = 1;
}