Add part 2 for net testing

2018/July/net-testing/net.go

@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"errors"
 	"fmt"
+	"log"
 	"net"
 	"strings"
 )
@@ -24,6 +25,9 @@ func NewServer(protocol, addr string) (Server, error) {
 			addr: addr,
 		}, nil
 	case "udp":
+		return &UDPServer{
+			addr: addr,
+		}, nil
 	}
 	return nil, errors.New("Invalid protocol given")
 }
@@ -97,3 +101,52 @@ func (t *TCPServer) handleConnection(conn net.Conn) {
 		rw.Flush()
 	}
 }
+
+// UDPServer holds the necessary structure for our
+// UDP server.
+type UDPServer struct {
+	addr   string
+	server *net.UDPConn
+}
+
+// Run starts the UDP server.
+func (u *UDPServer) Run() (err error) {
+	laddr, err := net.ResolveUDPAddr("udp", u.addr)
+	if err != nil {
+		return errors.New("could not resolve UDP addr")
+	}
+
+	u.server, err = net.ListenUDP("udp", laddr)
+	if err != nil {
+		return errors.New("could not listen on UDP")
+	}
+
+	return u.handleConnections()
+}
+
+func (u *UDPServer) handleConnections() error {
+	var err error
+	for {
+		buf := make([]byte, 2048)
+		n, conn, err := u.server.ReadFromUDP(buf)
+		if err != nil {
+			log.Println(err)
+			break
+		}
+		if conn == nil {
+			continue
+		}
+
+		go u.handleConnection(conn, buf[:n])
+	}
+	return err
+}
+
+func (u *UDPServer) handleConnection(addr *net.UDPAddr, cmd []byte) {
+	u.server.WriteToUDP([]byte(fmt.Sprintf("Request recieved: %s", cmd)), addr)
+}
+
+// Close ensures that the UDPServer is shut down gracefully.
+func (u *UDPServer) Close() error {
+	return u.server.Close()
+}

2018/July/net-testing/net_test.go

@@ -7,33 +7,59 @@ import (
 	"testing"
 )
 
-var srv Server
+var tcp, udp Server
 
 func init() {
 	// Start the new server
-	srv, err := NewServer("tcp", ":1123")
+	tcp, err := NewServer("tcp", ":1123")
 	if err != nil {
 		log.Println("error starting TCP server")
 		return
 	}
 
-	// Run the server in goroutine to stop blocking
+	udp, err := NewServer("udp", ":6250")
+	if err != nil {
+		log.Println("error starting UDP server")
+		return
+	}
+
+	// Run the servers in goroutines to stop blocking
 	go func() {
-		srv.Run()
+		tcp.Run()
+	}()
+	go func() {
+		udp.Run()
 	}()
 }
 
 func TestNETServer_Running(t *testing.T) {
 	// Simply check that the server is up and can
 	// accept connections.
-	conn, err := net.Dial("tcp", ":1123")
-	if err != nil {
-		t.Error("could not connect to server: ", err)
+	servers := []struct {
+		protocol string
+		addr     string
+	}{
+		{"tcp", ":1123"},
+		{"udp", ":6250"},
+	}
+	for _, serv := range servers {
+		conn, err := net.Dial(serv.protocol, serv.addr)
+		if err != nil {
+			t.Error("could not connect to server: ", err)
+		}
+		defer conn.Close()
 	}
-	defer conn.Close()
 }
 
 func TestNETServer_Request(t *testing.T) {
+	servers := []struct {
+		protocol string
+		addr     string
+	}{
+		{"tcp", ":1123"},
+		{"udp", ":6250"},
+	}
+
 	tt := []struct {
 		test    string
 		payload []byte
@@ -43,26 +69,28 @@ func TestNETServer_Request(t *testing.T) {
 		{"Sending another simple request works", []byte("goodbye world\n"), []byte("Request received: goodbye world")},
 	}
 
-	for _, tc := range tt {
-		t.Run(tc.test, func(t *testing.T) {
-			conn, err := net.Dial("tcp", ":1123")
-			if err != nil {
-				t.Error("could not connect to TCP server: ", err)
-			}
-			defer conn.Close()
+	for _, serv := range servers {
+		for _, tc := range tt {
+			t.Run(tc.test, func(t *testing.T) {
+				conn, err := net.Dial(serv.protocol, serv.addr)
+				if err != nil {
+					t.Error("could not connect to server: ", err)
+				}
+				defer conn.Close()
 
-			if _, err := conn.Write(tc.payload); err != nil {
-				t.Error("could not write payload to TCP server:", err)
-			}
+				if _, err := conn.Write(tc.payload); err != nil {
+					t.Error("could not write payload to server:", err)
+				}
 
-			out := make([]byte, 1024)
-			if _, err := conn.Read(out); err == nil {
-				if bytes.Compare(out, tc.want) == 0 {
-					t.Error("response did match expected output")
+				out := make([]byte, 1024)
+				if _, err := conn.Read(out); err == nil {
+					if bytes.Compare(out, tc.want) == 0 {
+						t.Error("response did match expected output")
+					}
+				} else {
+					t.Error("could not read from connection")
 				}
-			} else {
-				t.Error("could not read from connection")
-			}
-		})
+			})
+		}
 	}
 }