Improve test cleanup + fix RegisterSSHClientVersionPicker issue

psiphon/controller_test.go

@@ -44,19 +44,9 @@ import (
 	"github.com/elazarl/goproxy"
 )
 
-var testDataDirName string
-
 func TestMain(m *testing.M) {
 	flag.Parse()
 
-	var err error
-	testDataDirName, err = ioutil.TempDir("", "psiphon-controller-test")
-	if err != nil {
-		fmt.Printf("TempDir failed: %s\n", err)
-		os.Exit(1)
-	}
-	defer os.RemoveAll(testDataDirName)
-
 	SetEmitDiagnosticNotices(true)
 
 	initDisruptor()
@@ -462,6 +452,12 @@ type controllerRunConfig struct {
 
 func controllerRun(t *testing.T, runConfig *controllerRunConfig) {
 
+	testDataDirName, err := ioutil.TempDir("", "psiphon-controller-test")
+	if err != nil {
+		t.Fatalf("TempDir failed: %s\n", err)
+	}
+	defer os.RemoveAll(testDataDirName)
+
 	configJSON, err := ioutil.ReadFile("controller_test.config")
 	if err != nil {
 		// Skip, don't fail, if config file is not present
@@ -554,9 +550,6 @@ func controllerRun(t *testing.T, runConfig *controllerRunConfig) {
 	// that the tactics request succeeds.
 	config.NetworkID = "NETWORK1"
 
-	os.Remove(config.UpgradeDownloadFilename)
-	os.Remove(config.RemoteServerListDownloadFilename)
-
 	err = OpenDataStore(config)
 	if err != nil {
 		t.Fatalf("error initializing datastore: %s", err)

psiphon/dialParameters_test.go

@@ -55,14 +55,14 @@ func runDialParametersAndReplay(t *testing.T, tunnelProtocol string) {
 
 	t.Logf("Test %s...", tunnelProtocol)
 
-	SetNoticeWriter(ioutil.Discard)
-
 	testDataDirName, err := ioutil.TempDir("", "psiphon-dial-parameters-test")
 	if err != nil {
 		t.Fatalf("TempDir failed: %s", err)
 	}
 	defer os.RemoveAll(testDataDirName)
 
+	SetNoticeWriter(ioutil.Discard)
+
 	clientConfig := &Config{
 		PropagationChannelId: "0",
 		SponsorId:            "0",
@@ -100,7 +100,7 @@ func runDialParametersAndReplay(t *testing.T, tunnelProtocol string) {
 	}
 
 	RegisterSSHClientVersionPicker(func() string {
-		versions := []string{"v1", "v2", "v3"}
+		versions := []string{"SSH-2.0-A", "SSH-2.0-B", "SSH-2.0-C"}
 		return versions[prng.Intn(len(versions))]
 	})
 

psiphon/limitProtocols_test.go

@@ -36,6 +36,12 @@ import (
 
 func TestLimitTunnelProtocols(t *testing.T) {
 
+	testDataDirName, err := ioutil.TempDir("", "psiphon-limit-tunnel-protocols-test")
+	if err != nil {
+		t.Fatalf("TempDir failed: %s", err)
+	}
+	defer os.RemoveAll(testDataDirName)
+
 	initialLimitTunnelProtocols := protocol.TunnelProtocols{"OSSH", "UNFRONTED-MEEK-HTTPS-OSSH"}
 	initialLimitTunnelProtocolsCandidateCount := 100
 	limitTunnelProtocols := protocol.TunnelProtocols{"SSH", "UNFRONTED-MEEK-OSSH"}
@@ -88,12 +94,6 @@ func TestLimitTunnelProtocols(t *testing.T) {
 			}
 		}))
 
-	testDataDirName, err := ioutil.TempDir("", "psiphon-limit-tunnel-protocols-test")
-	if err != nil {
-		t.Fatalf("TempDir failed: %s", err)
-	}
-	defer os.RemoveAll(testDataDirName)
-
 	clientConfigJSON := `
     {
         "ClientPlatform" : "Windows",

psiphon/remoteServerList_test.go

@@ -250,6 +250,9 @@ func testObfuscatedRemoteServerLists(t *testing.T, omitMD5Sums bool) {
 
 	obfuscatedServerListRootURLsJSONConfig := "["
 	obfuscatedServerListRootURLs := make([]string, len(remoteServerListHostAddresses))
+
+	httpServers := make(chan *http.Server, len(remoteServerListHostAddresses))
+
 	for i := 0; i < len(remoteServerListHostAddresses); i++ {
 
 		obfuscatedServerListRootURLs[i] = fmt.Sprintf("http://%s/", remoteServerListHostAddresses[i])
@@ -278,10 +281,18 @@ func testObfuscatedRemoteServerLists(t *testing.T, omitMD5Sums bool) {
 				Addr:    remoteServerListHostAddress,
 				Handler: serveMux,
 			}
+			httpServers <- httpServer
 			httpServer.Serve(listener)
 		}(remoteServerListListeners[i], remoteServerListHostAddresses[i])
 	}
 
+	defer func() {
+		for i := 0; i < len(remoteServerListHostAddresses); i++ {
+			httpServer := <-httpServers
+			httpServer.Close()
+		}
+	}()
+
 	obfuscatedServerListDownloadDirectory := testDataDirName
 
 	//
@@ -366,6 +377,7 @@ func testObfuscatedRemoteServerLists(t *testing.T, omitMD5Sums bool) {
 	// Note: calling LoadConfig ensures all *int config fields are initialized
 	clientConfigJSONTemplate := `
     {
+        "DataStoreDirectory" : "%s",
         "ClientPlatform" : "",
         "ClientVersion" : "0",
         "SponsorId" : "0",
@@ -383,6 +395,7 @@ func testObfuscatedRemoteServerLists(t *testing.T, omitMD5Sums bool) {
 
 	clientConfigJSON := fmt.Sprintf(
 		clientConfigJSONTemplate,
+		testDataDirName,
 		signingPublicKey,
 		remoteServerListURL,
 		remoteServerListDownloadFilename,

psiphon/userAgent_test.go

@@ -22,8 +22,10 @@ package psiphon
 import (
 	"context"
 	"fmt"
+	"io/ioutil"
 	"net"
 	"net/http"
+	"os"
 	"sync"
 	"testing"
 	"time"
@@ -151,13 +153,18 @@ func initUserAgentCounterUpstreamProxy() {
 func attemptConnectionsWithUserAgent(
 	t *testing.T, tunnelProtocol string, isCONNECT bool) {
 
+	testDataDirName, err := ioutil.TempDir("", "psiphon-user-agent-test")
+	if err != nil {
+		t.Fatalf("TempDir failed: %s\n", err)
+	}
+	defer os.RemoveAll(testDataDirName)
+
 	initMockUserAgentPicker()
 	initUserAgentCounterUpstreamProxy()
 	resetUserAgentCounts()
 
 	// create a server entry
 
-	var err error
 	serverIPaddress := ""
 	for _, interfaceName := range []string{"eth0", "en0"} {
 		var serverIPv4Address, serverIPv6Address net.IP