<commit_before>package network

import (
	"context"
	"crypto"
	"net"
	"testing"
	"time"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow/networking/benchlist"
	"github.com/ava-labs/avalanchego/snow/validators"
	"github.com/ava-labs/avalanchego/utils"
	"github.com/ava-labs/avalanchego/utils/hashing"
	"github.com/ava-labs/avalanchego/utils/logging"
	"github.com/ava-labs/avalanchego/version"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/stretchr/testify/assert"
)

type TestMsg struct {
	op    Op
	bytes []byte
}

func newTestMsg(op Op, bits []byte) *TestMsg {
	return &TestMsg{op: op, bytes: bits}
}

func (m *TestMsg) Op() Op {
	return m.op
}

func (*TestMsg) Get(Field) interface{} {
	return nil
}

func (m *TestMsg) Bytes() []byte {
	return m.bytes
}

func TestPeer_Close(t *testing.T) {
	log := logging.NoLog{}
	ip := utils.NewDynamicIPDesc(
		net.IPv6loopback,
		0,
	)
	id := ids.ShortID(hashing.ComputeHash160Array([]byte(ip.IP().String())))
	networkID := uint32(0)
	appVersion := version.NewDefaultApplication("app", 0, 1, 0)
	versionParser := version.NewDefaultApplicationParser()

	listener := &testListener{
		addr: &net.TCPAddr{
			IP:   net.IPv6loopback,
			Port: 0,
		},
		inbound: make(chan net.Conn, 1<<10),
		closed:  make(chan struct{}),
	}
	caller := &testDialer{
		addr: &net.TCPAddr{
			IP:   net.IPv6loopback,
			Port: 0,
		},
		outbounds: make(map[string]*testListener),
	}
	serverUpgrader0 := NewTLSServerUpgrader(tlsConfig0)
	clientUpgrader0 := NewTLSClientUpgrader(tlsConfig0)

	vdrs := validators.NewSet()
	handler := &testHandler{}

	versionManager := version.NewCompatibility(
		appVersion,
		appVersion,
		time.Now(),
		appVersion,
		appVersion,
		time.Now(),
		appVersion,
	)

	netwrk := NewDefaultNetwork(
		prometheus.NewRegistry(),
		log,
		id,
		ip,
		networkID,
		versionManager,
		versionParser,
		listener,
		caller,
		serverUpgrader0,
		clientUpgrader0,
		vdrs,
		vdrs,
		handler,
		time.Duration(0),
		0,
		defaultSendQueueSize,
		HealthConfig{},
		benchlist.NewManager(&benchlist.Config{}),
		defaultAliasTimeout,
		cert0.PrivateKey.(crypto.Signer),
		defaultPeerListSize,
		defaultGossipPeerListTo,
		defaultGossipPeerListFreq,
		NewDialerConfig(0, 30*time.Second),
		false,
		defaultGossipAcceptedFrontierSize,
		defaultGossipOnAcceptSize,
	)
	assert.NotNil(t, netwrk)

	ip1 := utils.NewDynamicIPDesc(
		net.IPv6loopback,
		1,
	)
	caller.outbounds[ip1.IP().String()] = listener
	conn, err := caller.Dial(context.Background(), ip1.IP())
	assert.NoError(t, err)

	basenetwork := netwrk.(*network)

	newmsgbytes := []byte("hello")

	// fake a peer, and write a message
	peer := newPeer(basenetwork, conn, ip1.IP())
	peer.sender = make(chan []byte, 10)
	testMsg := newTestMsg(GetVersion, newmsgbytes)
	peer.Send(testMsg, true)

	// make sure the net pending and peer pending bytes updated
	if basenetwork.pendingBytes != int64(len(newmsgbytes)) {
		t.Fatalf("pending bytes invalid")
	}
	if peer.pendingBytes != int64(len(newmsgbytes)) {
		t.Fatalf("pending bytes invalid")
	}

	go func() {
		err := netwrk.Close()
		assert.NoError(t, err)
	}()

	peer.Close()

	// The network pending bytes should be reduced back to zero on close.
	if basenetwork.pendingBytes != int64(0) {
		t.Fatalf("pending bytes invalid")
	}
	if peer.pendingBytes != int64(len(newmsgbytes)) {
		t.Fatalf("pending bytes invalid")
	}
}<commit_msg>add test<commit_after>package network

import (
	"context"
	"crypto"
	"net"
	"testing"
	"time"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow/networking/benchlist"
	"github.com/ava-labs/avalanchego/snow/validators"
	"github.com/ava-labs/avalanchego/utils"
	"github.com/ava-labs/avalanchego/utils/hashing"
	"github.com/ava-labs/avalanchego/utils/logging"
	"github.com/ava-labs/avalanchego/version"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/stretchr/testify/assert"
)

type TestMsg struct {
	op    Op
	bytes []byte
}

func newTestMsg(op Op, bits []byte) *TestMsg {
	return &TestMsg{op: op, bytes: bits}
}

func (m *TestMsg) Op() Op {
	return m.op
}

func (*TestMsg) Get(Field) interface{} {
	return nil
}

func (m *TestMsg) Bytes() []byte {
	return m.bytes
}

func TestPeer_Close(t *testing.T) {
	log := logging.NoLog{}
	ip := utils.NewDynamicIPDesc(
		net.IPv6loopback,
		0,
	)
	id := ids.ShortID(hashing.ComputeHash160Array([]byte(ip.IP().String())))
	networkID := uint32(0)
	appVersion := version.NewDefaultApplication("app", 0, 1, 0)
	versionParser := version.NewDefaultApplicationParser()

	listener := &testListener{
		addr: &net.TCPAddr{
			IP:   net.IPv6loopback,
			Port: 0,
		},
		inbound: make(chan net.Conn, 1<<10),
		closed:  make(chan struct{}),
	}
	caller := &testDialer{
		addr: &net.TCPAddr{
			IP:   net.IPv6loopback,
			Port: 0,
		},
		outbounds: make(map[string]*testListener),
	}
	serverUpgrader0 := NewTLSServerUpgrader(tlsConfig0)
	clientUpgrader0 := NewTLSClientUpgrader(tlsConfig0)

	vdrs := validators.NewSet()
	handler := &testHandler{}

	versionManager := version.NewCompatibility(
		appVersion,
		appVersion,
		time.Now(),
		appVersion,
		appVersion,
		time.Now(),
		appVersion,
	)

	netwrk := NewDefaultNetwork(
		prometheus.NewRegistry(),
		log,
		id,
		ip,
		networkID,
		versionManager,
		versionParser,
		listener,
		caller,
		serverUpgrader0,
		clientUpgrader0,
		vdrs,
		vdrs,
		handler,
		time.Duration(0),
		0,
		defaultSendQueueSize,
		HealthConfig{},
		benchlist.NewManager(&benchlist.Config{}),
		defaultAliasTimeout,
		cert0.PrivateKey.(crypto.Signer),
		defaultPeerListSize,
		defaultGossipPeerListTo,
		defaultGossipPeerListFreq,
		NewDialerConfig(0, 30*time.Second),
		false,
		defaultGossipAcceptedFrontierSize,
		defaultGossipOnAcceptSize,
	)
	assert.NotNil(t, netwrk)

	ip1 := utils.NewDynamicIPDesc(
		net.IPv6loopback,
		1,
	)
	caller.outbounds[ip1.IP().String()] = listener
	conn, err := caller.Dial(context.Background(), ip1.IP())
	assert.NoError(t, err)

	basenetwork := netwrk.(*network)

	newmsgbytes := []byte("hello")

	// fake a peer, and write a message
	peer := newPeer(basenetwork, conn, ip1.IP())
	peer.sender = make(chan []byte, 10)
	testMsg := newTestMsg(GetVersion, newmsgbytes)
	peer.Send(testMsg, true)

	// make sure the net pending and peer pending bytes updated
	if basenetwork.pendingBytes != int64(len(newmsgbytes)) {
		t.Fatalf("pending bytes invalid")
	}
	if peer.pendingBytes != int64(len(newmsgbytes)) {
		t.Fatalf("pending bytes invalid")
	}

	go func() {
		err := netwrk.Close()
		assert.NoError(t, err)
	}()

	peer.Close()

	// The network pending bytes should be reduced back to zero on close.
	if basenetwork.pendingBytes != int64(0) {
		t.Fatalf("pending bytes invalid")
	}
	if peer.pendingBytes != int64(len(newmsgbytes)) {
		t.Fatalf("pending bytes invalid")
	}
}

func TestShouldCompress(t *testing.T) {
	p := newPeer(nil, nil, utils.IPDesc{}) // Dummy values are OK here
	shouldNotCompressOps := []Op{Version, GetVersion, PeerList, GetPeerList, Ping, Pong}
	shouldCompressOps := []Op{Put, Get, MultiPut, GetAncestors, Accepted, AcceptedFrontier, GetAccepted, GetAcceptedFrontier, PushQuery, PullQuery, Chits}
	for _, op := range shouldCompressOps {
		// Shouldn't compress because [p.canHandleCompressed] == false
		assert.False(t, p.shouldCompress(op))
	}
	for _, op := range shouldNotCompressOps {
		assert.False(t, p.shouldCompress(op))
	}

	p.canHandleCompressed = true
	for _, op := range shouldCompressOps {
		assert.True(t, p.shouldCompress(op))
	}
	for _, op := range shouldNotCompressOps {
		assert.False(t, p.shouldCompress(op))
	}
}
{'fraction_non_alphanumeric': 0.14163259923513122, 'fraction_numerical': 0.01239614928128709, 'mean_word_length': 4.556043956043956, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2717127', 'n_tokens_mistral': 3046, 'n_tokens_neox': 2854, 'n_words': 541}