From 4549996a9b9ccee3ed993010a6d5f1c1e1d4913f Mon Sep 17 00:00:00 2001
From: Rafael Sampaio <5679073+r4f4ss@users.noreply.github.com>
Date: Fri, 4 Oct 2024 18:50:39 -0300
Subject: [PATCH] add test for function doPortMapping

---
 cmd/shisui/main_test.go | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)
 create mode 100644 cmd/shisui/main_test.go

diff --git a/cmd/shisui/main_test.go b/cmd/shisui/main_test.go
new file mode 100644
index 0000000000..1a09a9ae10
--- /dev/null
+++ b/cmd/shisui/main_test.go
@@ -0,0 +1,33 @@
+package main
+
+import (
+	"net"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/p2p/enode"
+	"github.com/ethereum/go-ethereum/p2p/nat"
+	"github.com/stretchr/testify/assert"
+)
+
+func newLocalNodeForTesting() (*enode.LocalNode, *enode.DB) {
+	db, _ := enode.OpenDB("")
+	key, _ := crypto.GenerateKey()
+	return enode.NewLocalNode(db, key), db
+}
+
+func TestDoPortMapping(t *testing.T) {
+	nat := nat.ExtIP{33, 44, 55, 66}
+	localNode, _ := newLocalNodeForTesting()
+	listenerAddr := &net.UDPAddr{IP: net.IP{127, 0, 0, 1}, Port: 1234}
+
+	doPortMapping(nat, localNode, listenerAddr)
+
+	assert.Equal(t, localNode.Seq(), uint64(1))
+	assert.Equal(t, localNode.Node().IP(), net.IP{33, 44, 55, 66})
+	assert.Equal(t, localNode.Node().UDP(), 1234)
+	assert.Equal(t, localNode.Node().TCP(), 0)
+
+	_ = localNode.Node().UDP()
+	assert.Equal(t, localNode.Seq(), uint64(2))
+}