package throttler

import (
	"encoding/binary"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/vishvananda/netlink"
)

func TestWithDstPortPort(t *testing.T) {
	sel := WithDstPort(nil, 123)
	assert.Equal(t, uint8(netlink.TC_U32_TERMINAL), sel.Flags)
	assert.Equal(t, uint8(0x1), sel.Nkeys)
	assert.Equal(t, 1, len(sel.Keys))
	assert.Equal(t, netlink.TcU32Key{
		Mask: binary.BigEndian.Uint32([]byte{0x00, 0x00, 0xff, 0xff}),
		Val:  binary.BigEndian.Uint32([]byte{0x00, 0x00, 0x00, 0x7b}), // buf[2:] = hex(123)
		Off:  ipHeaderBytes,
	}, sel.Keys[0])
}

func TestWithSrcPort(t *testing.T) {
	sel := WithSrcPort(nil, 456)
	assert.Equal(t, uint8(netlink.TC_U32_TERMINAL), sel.Flags)
	assert.Equal(t, uint8(0x1), sel.Nkeys)
	assert.Equal(t, 1, len(sel.Keys))
	assert.Equal(t, netlink.TcU32Key{
		Mask: binary.BigEndian.Uint32([]byte{0xff, 0xff, 0x00, 0x00}),
		Val:  binary.BigEndian.Uint32([]byte{0x01, 0xc8, 0x00, 0x00}), // buf[:2] = hex(456)
		Off:  ipHeaderBytes,
	}, sel.Keys[0])
}