1 files changed, 87 insertions, 0 deletions
diff --git a/table/message_test.go b/table/message_test.go
index 01c9ae58..581c6aba 100644
--- a/table/message_test.go
+++ b/table/message_test.go
@@ -18,6 +18,7 @@ package table
 import (
 	"github.com/osrg/gobgp/packet/bgp"
 	"github.com/stretchr/testify/assert"
+	"reflect"
 	"testing"
 	"time"
 )
@@ -324,6 +325,92 @@ func TestAsPathAs4TransInvalid4(t *testing.T) {
 	assert.Equal(t, msg.PathAttributes[0].(*bgp.PathAttributeAsPath).Value[0].(*bgp.As4PathParam).AS[4], uint32(40001))
 }
 
+func TestAggregator4BytesASes(t *testing.T) {
+	getAggr := func(msg *bgp.BGPUpdate) *bgp.PathAttributeAggregator {
+		for _, attr := range msg.PathAttributes {
+			switch attr.(type) {
+			case *bgp.PathAttributeAggregator:
+				return attr.(*bgp.PathAttributeAggregator)
+			}
+		}
+		return nil
+	}
+
+	getAggr4 := func(msg *bgp.BGPUpdate) *bgp.PathAttributeAs4Aggregator {
+		for _, attr := range msg.PathAttributes {
+			switch attr.(type) {
+			case *bgp.PathAttributeAs4Aggregator:
+				return attr.(*bgp.PathAttributeAs4Aggregator)
+			}
+		}
+		return nil
+	}
+
+	addr := "192.168.0.1"
+	as4 := uint32(100000)
+	as := uint32(1000)
+	msg := bgp.NewBGPUpdateMessage(nil, []bgp.PathAttributeInterface{bgp.NewPathAttributeAggregator(as4, addr)}, nil).Body.(*bgp.BGPUpdate)
+
+	// 4byte capable to 4byte capable for 4 bytes AS
+	assert.Equal(t, UpdatePathAggregator4ByteAs(msg), nil)
+	assert.Equal(t, getAggr(msg).Value.AS, as4)
+	assert.Equal(t, getAggr(msg).Value.Address.String(), addr)
+
+	// 4byte capable to 2byte capable for 4 bytes AS
+	UpdatePathAggregator2ByteAs(msg)
+	assert.Equal(t, getAggr(msg).Value.AS, uint32(bgp.AS_TRANS))
+	assert.Equal(t, getAggr(msg).Value.Askind, reflect.Uint16)
+	assert.Equal(t, getAggr4(msg).Value.AS, as4)
+	assert.Equal(t, getAggr4(msg).Value.Address.String(), addr)
+
+	msg = bgp.NewBGPUpdateMessage(nil, []bgp.PathAttributeInterface{bgp.NewPathAttributeAggregator(uint16(bgp.AS_TRANS), addr), bgp.NewPathAttributeAs4Aggregator(as4, addr)}, nil).Body.(*bgp.BGPUpdate)
+	assert.Equal(t, getAggr(msg).Value.AS, uint32(bgp.AS_TRANS))
+	assert.Equal(t, getAggr(msg).Value.Askind, reflect.Uint16)
+
+	// non 4byte capable to 4byte capable for 4 bytes AS
+	assert.Equal(t, UpdatePathAggregator4ByteAs(msg), nil)
+	assert.Equal(t, getAggr(msg).Value.AS, as4)
+	assert.Equal(t, getAggr(msg).Value.Askind, reflect.Uint32)
+	assert.Equal(t, getAggr(msg).Value.Address.String(), addr)
+	assert.Equal(t, getAggr4(msg), (*bgp.PathAttributeAs4Aggregator)(nil))
+
+	// non 4byte capable to non 4byte capable for 4 bytes AS
+	UpdatePathAggregator2ByteAs(msg)
+	assert.Equal(t, getAggr(msg).Value.AS, uint32(bgp.AS_TRANS))
+	assert.Equal(t, getAggr(msg).Value.Askind, reflect.Uint16)
+	assert.Equal(t, getAggr4(msg).Value.AS, as4)
+	assert.Equal(t, getAggr4(msg).Value.Address.String(), addr)
+
+	msg = bgp.NewBGPUpdateMessage(nil, []bgp.PathAttributeInterface{bgp.NewPathAttributeAggregator(uint32(as), addr)}, nil).Body.(*bgp.BGPUpdate)
+	// 4byte capable to 4byte capable for 2 bytes AS
+	assert.Equal(t, getAggr(msg).Value.AS, as)
+	assert.Equal(t, getAggr(msg).Value.Askind, reflect.Uint32)
+	assert.Equal(t, UpdatePathAggregator4ByteAs(msg), nil)
+	assert.Equal(t, getAggr(msg).Value.AS, as)
+	assert.Equal(t, getAggr(msg).Value.Askind, reflect.Uint32)
+
+	// 4byte capable to non 4byte capable for 2 bytes AS
+	UpdatePathAggregator2ByteAs(msg)
+	assert.Equal(t, getAggr4(msg), (*bgp.PathAttributeAs4Aggregator)(nil))
+	assert.Equal(t, getAggr(msg).Value.Askind, reflect.Uint16)
+	assert.Equal(t, getAggr(msg).Value.AS, as)
+
+	msg = bgp.NewBGPUpdateMessage(nil, []bgp.PathAttributeInterface{bgp.NewPathAttributeAggregator(uint16(as), addr)}, nil).Body.(*bgp.BGPUpdate)
+	// non 4byte capable to 4byte capable for 2 bytes AS
+	assert.Equal(t, getAggr(msg).Value.AS, as)
+	assert.Equal(t, getAggr(msg).Value.Askind, reflect.Uint16)
+	assert.Equal(t, UpdatePathAggregator4ByteAs(msg), nil)
+
+	assert.Equal(t, getAggr(msg).Value.AS, as)
+	assert.Equal(t, getAggr(msg).Value.Askind, reflect.Uint32)
+
+	// non 4byte capable to non 4byte capable for 2 bytes AS
+	UpdatePathAggregator2ByteAs(msg)
+	assert.Equal(t, getAggr(msg).Value.AS, as)
+	assert.Equal(t, getAggr(msg).Value.Askind, reflect.Uint16)
+	assert.Equal(t, getAggr4(msg), (*bgp.PathAttributeAs4Aggregator)(nil))
+}
+
 func TestBMP(t *testing.T) {
 	aspath1 := []bgp.AsPathParamInterface{
 		bgp.NewAs4PathParam(2, []uint32{1000000}),