From aab2d5ba841461753bd4c6c14c0e9cafe85039e1 Mon Sep 17 00:00:00 2001
From: Hugo Arregui <hugo@decentraland.org>
Date: Tue, 13 Aug 2019 14:49:36 -0300
Subject: [PATCH] Use vnet for ConnectivityOnStartup test

The test was failing if the first ContactCandidates was not enough to
establish the connection.
---
 agent_test.go | 87 ++++++++++++++++++++++++++++++++++++---------------
 1 file changed, 61 insertions(+), 26 deletions(-)

diff --git a/agent_test.go b/agent_test.go
index b5c3736..93c130a 100644
--- a/agent_test.go
+++ b/agent_test.go
@@ -7,9 +7,10 @@ import (
 	"testing"
 	"time"
 
-	"github.com/pion/logging"
 	"github.com/pion/stun"
 	"github.com/pion/transport/test"
+	"github.com/pion/transport/vnet"
+	"github.com/stretchr/testify/require"
 )
 
 type mockPacketConn struct {
@@ -356,45 +357,79 @@ func TestHandlePeerReflexive(t *testing.T) {
 	})
 }
 
-// Assert that Agent on startup sends message, and doesn't wait for taskloop to start
+// Assert that Agent on startup sends message, and doesn't wait for connectivityTicker to fire
 // github.com/pion/ice/issues/15
 func TestConnectivityOnStartup(t *testing.T) {
-	lim := test.TimeOut(time.Second * 5)
-	defer lim.Stop()
-
-	cfg := &AgentConfig{
-		Urls:             []*URL{},
-		Trickle:          false,
-		NetworkTypes:     supportedNetworkTypes,
-		taskLoopInterval: time.Hour,
-		LoggerFactory:    logging.NewDefaultLoggerFactory(),
+	stunServerURL := &URL{
+		Scheme: SchemeTypeSTUN,
+		Host:   "1.2.3.4",
+		Port:   3478,
+		Proto:  ProtoTypeUDP,
 	}
 
+	v, err := buildVNet(&vnet.NATType{
+		MappingBehavior:   vnet.EndpointIndependent,
+		FilteringBehavior: vnet.EndpointIndependent,
+	})
+	require.NoError(t, err, "should succeed")
+	defer v.close()
+
 	aNotifier, aConnected := onConnected()
 	bNotifier, bConnected := onConnected()
 
-	aAgent, err := NewAgent(cfg)
-	if err != nil {
-		t.Error(err)
-	}
-	err = aAgent.OnConnectionStateChange(aNotifier)
-	if err != nil {
-		panic(err)
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	cfg0 := &AgentConfig{
+		Urls:             []*URL{stunServerURL},
+		Trickle:          true,
+		NetworkTypes:     supportedNetworkTypes,
+		MulticastDNSMode: MulticastDNSModeDisabled,
+		Net:              v.net0,
+
+		taskLoopInterval: time.Hour,
 	}
 
-	bAgent, err := NewAgent(cfg)
-	if err != nil {
-		t.Error(err)
-	}
-	err = bAgent.OnConnectionStateChange(bNotifier)
-	if err != nil {
-		panic(err)
+	aAgent, err := NewAgent(cfg0)
+	require.NoError(t, err)
+	require.NoError(t, aAgent.OnConnectionStateChange(aNotifier))
+	require.NoError(t, aAgent.OnCandidate(func(candidate Candidate) {
+		if candidate == nil {
+			wg.Done()
+		}
+	}))
+	require.NoError(t, aAgent.GatherCandidates())
+
+	cfg1 := &AgentConfig{
+		Urls:             []*URL{stunServerURL},
+		Trickle:          true,
+		NetworkTypes:     supportedNetworkTypes,
+		MulticastDNSMode: MulticastDNSModeDisabled,
+		Net:              v.net1,
+		taskLoopInterval: time.Hour,
 	}
 
-	connect(aAgent, bAgent)
+	bAgent, err := NewAgent(cfg1)
+	require.NoError(t, err)
+	require.NoError(t, bAgent.OnConnectionStateChange(bNotifier))
+	require.NoError(t, bAgent.OnCandidate(func(candidate Candidate) {
+		if candidate == nil {
+			wg.Done()
+		}
+	}))
+	require.NoError(t, bAgent.GatherCandidates())
 
+	wg.Wait()
+	aConn, bConn := connectWithVNet(aAgent, bAgent)
+
+	// Ensure pair selected
+	// Note: this assumes ConnectionStateConnected is thrown after selecting the final pair
 	<-aConnected
 	<-bConnected
+
+	if !closePipe(t, aConn, bConn) {
+		return
+	}
 }
 
 func TestInboundValidity(t *testing.T) {