Test UDPPing country parsing.

server/federation.go

@@ -73,7 +73,7 @@ func (s *Server) getNumSubs() int64 {
 // pings it, so we can determine our own external IP address.
 func (s *Server) getAndSetExternalIp(msg *pb.ServerMessage) error {
 	log.Println(msg)
-	myIp, err := UDPPing(msg.Address, msg.Port)
+	myIp, _, err := UDPPing(msg.Address, msg.Port)
 	if err != nil {
 		return err
 	}

server/udp.go

@@ -149,47 +149,52 @@ func (pong *SPVPong) DecodeAddress() net.IP {
 	)
 }
 
+func (pong *SPVPong) DecodeCountry() string {
+	return pb.Location_Country_name[int32(pong.country)]
+}
+
 // UDPPing sends a ping over udp to another hub and returns the ip address of
 // this hub.
-func UDPPing(ip, port string) (net.IP, error) {
+func UDPPing(ip, port string) (net.IP, string, error) {
 	address := ip + ":" + port
 	addr, err := net.ResolveUDPAddr("udp", address)
 	if err != nil {
-		return net.IP{}, err
+		return net.IP{}, "", err
 	}
 
 	conn, err := net.DialUDP("udp", nil, addr)
 	if err != nil {
-		return net.IP{}, err
+		return net.IP{}, "", err
 	}
 
 	defer conn.Close()
 
 	_, err = conn.Write(encodeSPVPing())
 	if err != nil {
-		return net.IP{}, err
+		return net.IP{}, "", err
 	}
 
 	buffer := make([]byte, maxBufferSize)
 	deadline := time.Now().Add(time.Second)
 	err = conn.SetReadDeadline(deadline)
 	if err != nil {
-		return net.IP{}, err
+		return net.IP{}, "", err
 	}
 	n, _, err := conn.ReadFromUDP(buffer)
 	if err != nil {
-		return net.IP{}, err
+		return net.IP{}, "", err
 	}
 
 	pong := decodeSPVPong(buffer[:n])
 
 	if pong == nil {
-		return net.IP{}, errors.Base("Pong decoding failed")
+		return net.IP{}, "", errors.Base("Pong decoding failed")
 	}
 
 	myAddr := pong.DecodeAddress()
+	country := pong.DecodeCountry()
 
-	return myAddr, nil
+	return myAddr, country, nil
 }
 
 // UDPServer is a goroutine that starts an udp server that implements the hubs

server/udp_test.go

@@ -13,12 +13,14 @@ func TestUDPPing(t *testing.T) {
 	args.StartUDP = false
 
 	tests := []struct {
-		name string
-		want string
+		name 		string
+		wantIP 		string
+		wantCountry string
 	} {
 		{
-			name: "Get the right ip from production server.",
-			want: "SETME",
+			name: 		 "Get the right ip from production server.",
+			wantIP: 	 "SETME",
+			wantCountry: "US",
 		},
 	}
 
@@ -28,7 +30,8 @@ func TestUDPPing(t *testing.T) {
 			toAddr := "spv16.lbry.com"
 			toPort := "50001"
 
-			ip, err := UDPPing(toAddr, toPort)
+			ip, country, err := UDPPing(toAddr, toPort)
+			gotCountry := country
 			if err != nil {
 				log.Println(err)
 			}
@@ -41,11 +44,14 @@ func TestUDPPing(t *testing.T) {
 
 			digIP := strings.TrimSpace(string(res))
 			udpIP := ip.String()
-			tt.want = digIP
+			tt.wantIP = digIP
 
-			got1 := udpIP
-			if got1 != tt.want {
-				t.Errorf("got: '%s', want: '%s'\n", got1, tt.want)
+			gotIP := udpIP
+			if gotIP != tt.wantIP {
+				t.Errorf("got: '%s', want: '%s'\n", gotIP, tt.wantIP)
+			}
+			if gotCountry != tt.wantCountry {
+				t.Errorf("got: '%s', want: '%s'\n", gotCountry, tt.wantCountry)
 			}
 		})
 	}