Fix crash for countries command (NordSecurity#108)

Signed-off-by: Marius Sincovici <[email protected]>

cli/cli_connect_test.go

@@ -4,8 +4,9 @@ import (
 	"bytes"
 	"context"
 	"flag"
-	"log"
+	"io"
 	"os"
+	"strings"
 	"testing"
 
 	"github.com/NordSecurity/nordvpn-linux/client/config"
@@ -18,39 +19,87 @@ import (
 	"google.golang.org/grpc"
 )
 
+func captureOutput(f func()) (string, error) {
+	reader, writer, err := os.Pipe()
+	if err != nil {
+		return "", err
+	}
+	stdout := os.Stdout
+	stderr := os.Stderr
+	defer func() {
+		os.Stdout = stdout
+		os.Stderr = stderr
+	}()
+
+	os.Stdout = writer
+	os.Stderr = writer
+
+	f()
+
+	writer.Close() // close to unblock io.Copy(&buf, reader)
+	var buf bytes.Buffer
+	io.Copy(&buf, reader)
+	return strings.TrimSuffix(buf.String(), "\n"), nil
+}
+
 type mockDaemonClient struct {
 	pb.DaemonClient
+	cities    []string
+	groups    []string
+	countries []string
 }
 
 func (c mockDaemonClient) Cities(ctx context.Context, in *pb.CitiesRequest, opts ...grpc.CallOption) (*pb.Payload, error) {
-	x := &pb.Payload{
-		Type: internal.CodeSuccess,
-		Data: []string{"Paris", "Madrid", "Atlanta", "Chicago", "Los_Angeles", "Miami", "New_York"},
+	if c.cities != nil {
+		return &pb.Payload{
+			Type: internal.CodeSuccess,
+			Data: c.cities,
+		}, nil
+	} else {
+		return &pb.Payload{
+			Type: internal.CodeEmptyPayloadError,
+			Data: nil,
+		}, nil
 	}
-	return x, nil
 }
+
 func (c mockDaemonClient) Countries(ctx context.Context, in *pb.CountriesRequest, opts ...grpc.CallOption) (*pb.Payload, error) {
-	x := &pb.Payload{
-		Type: internal.CodeSuccess,
-		Data: []string{"Canada", "France", "Germany", "Hong_Kong", "Italy", "Japan", "Netherlands", "Poland", "Singapore", "Spain", "Sweden", "Switzerland", "Spain", "Turkey", "United_Arab_Emirates", "United_Kingdom", "United_States"},
+	if c.countries != nil {
+		return &pb.Payload{
+			Type: internal.CodeSuccess,
+			Data: c.countries,
+		}, nil
+	} else {
+		return &pb.Payload{
+			Type: internal.CodeEmptyPayloadError,
+			Data: nil,
+		}, nil
 	}
-	return x, nil
 }
 func (c mockDaemonClient) Groups(ctx context.Context, in *pb.GroupsRequest, opts ...grpc.CallOption) (*pb.Payload, error) {
-	x := &pb.Payload{
-		Type: internal.CodeSuccess,
-		Data: []string{"Africa_The_Middle_East_And_India", "Asia_Pacific", "Europe", "Obfuscated_Servers", "The_Americas"},
+	if c.groups != nil {
+		return &pb.Payload{
+			Type: internal.CodeSuccess,
+			Data: c.groups,
+		}, nil
+	} else {
+		return &pb.Payload{
+			Type: internal.CodeEmptyPayloadError,
+			Data: nil,
+		}, nil
 	}
-	return x, nil
 }
 
 func TestConnectAutoComplete(t *testing.T) {
 	category.Set(t, category.Unit)
-	c := cmd{mockDaemonClient{}, nil, nil, "", nil, config.Config{}, nil}
+	mockClient := mockDaemonClient{}
+	c := cmd{&mockClient, nil, nil, "", nil, config.Config{}, nil}
 	tests := []struct {
-		name     string
-		expected []string
-		input    []string
+		name      string
+		countries []string
+		groups    []string
+		expected  []string
+		input     []string
 	}{
 		{
 			name:     "France",
@@ -67,25 +116,34 @@ func TestConnectAutoComplete(t *testing.T) {
 			expected: []string{"Atlanta", "Chicago", "Los_Angeles", "Miami", "New_York"},
 			input:    []string{"United_States"},
 		},
-		{
-			name:     "Groups and Countries",
-			expected: []string{"Africa_The_Middle_East_And_India", "Asia_Pacific", "Europe", "Obfuscated_Servers", "The_Americas", "Canada", "France", "Germany", "Hong_Kong", "Italy", "Japan", "Netherlands", "Poland", "Singapore", "Spain", "Sweden", "Switzerland", "Spain", "Turkey", "United_Arab_Emirates", "United_Kingdom", "United_States"},
-			input:    []string{},
+		{ // in this case because input is empty, countries and groups will be displayed
+			name:      "Groups and Countries",
+			groups:    []string{"Europe", "Obfuscated_Servers", "The_Americas"},
+			countries: []string{"Canada", "France", "Germany"},
+			expected:  []string{"Canada", "France", "Germany", "Europe", "Obfuscated_Servers", "The_Americas"},
+			input:     []string{},
 		},
 	}
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			app := cli.NewApp()
 			set := flag.NewFlagSet("test", 0)
+			mockClient.cities = test.expected
+			mockClient.countries = test.countries
+			mockClient.groups = test.groups
 			set.Parse(test.input)
 			ctx := cli.NewContext(app, set, &cli.Context{Context: context.Background()})
-			var output bytes.Buffer
-			c.ConnectAutoComplete(ctx)
-			log.SetOutput(&output)
-			defer log.SetOutput(os.Stdout)
+
+			result, err := captureOutput(func() {
+				c.ConnectAutoComplete(ctx)
+			})
+
+			assert.Nil(t, err)
+
 			list, _ := internal.Columns(test.expected)
-			assert.Contains(t, output.String(), list)
+			assert.NotEmpty(t, list)
+			assert.Equal(t, list, result)
 		})
 	}
 }

cli/cli_countries.go

@@ -20,6 +20,13 @@ func (c *cmd) Countries(ctx *cli.Context) error {
 		Obfuscate: c.config.Obfuscate,
 	})
 	if err != nil {
+		log.Println(internal.ErrorPrefix, err)
+		return formatError(err)
+	}
+
+	if resp.Type != internal.CodeSuccess {
+		err := fmt.Errorf(MsgListIsEmpty, "countries")
+		log.Println(internal.ErrorPrefix, err)
 		return formatError(err)
 	}
 

cli/cli_countries_test.go

@@ -0,0 +1,53 @@
+package cli
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"testing"
+
+	"github.com/NordSecurity/nordvpn-linux/client/config"
+	"github.com/NordSecurity/nordvpn-linux/test/category"
+	"github.com/stretchr/testify/assert"
+	"github.com/urfave/cli/v2"
+)
+
+func TestCountriesList(t *testing.T) {
+	category.Set(t, category.Unit)
+	mockClient := mockDaemonClient{}
+	c := cmd{&mockClient, nil, nil, "", nil, config.Config{}, nil}
+
+	tests := []struct {
+		name          string
+		countries     []string
+		expected      string
+		input         string
+		expectedError error
+	}{
+		{
+			name:          "error response",
+			expectedError: formatError(fmt.Errorf(MsgListIsEmpty, "countries")),
+		},
+		{
+			name:      "countries list",
+			expected:  "France, Germany",
+			countries: []string{"France", "Germany"},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			app := cli.NewApp()
+			set := flag.NewFlagSet("test", 0)
+			mockClient.countries = test.countries
+			ctx := cli.NewContext(app, set, &cli.Context{Context: context.Background()})
+
+			result, err := captureOutput(func() {
+				err := c.Countries(ctx)
+				assert.Equal(t, test.expectedError, err)
+			})
+			assert.Nil(t, err)
+			assert.Equal(t, test.expected, result)
+		})
+	}
+}

cli/cli_groups_test.go

@@ -0,0 +1,53 @@
+package cli
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"testing"
+
+	"github.com/NordSecurity/nordvpn-linux/client/config"
+	"github.com/NordSecurity/nordvpn-linux/test/category"
+	"github.com/stretchr/testify/assert"
+	"github.com/urfave/cli/v2"
+)
+
+func TestGroupsList(t *testing.T) {
+	category.Set(t, category.Unit)
+	mockClient := mockDaemonClient{}
+	c := cmd{&mockClient, nil, nil, "", nil, config.Config{}, nil}
+
+	tests := []struct {
+		name          string
+		groups        []string
+		expected      string
+		input         string
+		expectedError error
+	}{
+		{
+			name:          "error response",
+			expectedError: formatError(fmt.Errorf(MsgListIsEmpty, "server groups")),
+		},
+		{
+			name:     "groups list",
+			expected: "group1, group2",
+			groups:   []string{"group1", "group2"},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			app := cli.NewApp()
+			set := flag.NewFlagSet("test", 0)
+			mockClient.groups = test.groups
+			ctx := cli.NewContext(app, set, &cli.Context{Context: context.Background()})
+
+			result, err := captureOutput(func() {
+				err := c.Groups(ctx)
+				assert.Equal(t, test.expectedError, err)
+			})
+			assert.Nil(t, err)
+			assert.Equal(t, test.expected, result)
+		})
+	}
+}

daemon/rpc_cities.go

@@ -16,6 +16,9 @@ func (r *RPC) Cities(ctx context.Context, in *pb.CitiesRequest) (*pb.Payload, er
 	var cfg config.Config
 	if err := r.cm.Load(&cfg); err != nil {
 		log.Println(internal.ErrorPrefix, err)
+		return &pb.Payload{
+			Type: internal.CodeConfigError,
+		}, nil
 	}
 
 	// collect cities and sort them

daemon/rpc_cities_test.go

@@ -0,0 +1,100 @@
+package daemon
+
+import (
+	"context"
+	"testing"
+
+	"github.com/NordSecurity/nordvpn-linux/config"
+	"github.com/NordSecurity/nordvpn-linux/daemon/pb"
+	"github.com/NordSecurity/nordvpn-linux/events/subs"
+	"github.com/NordSecurity/nordvpn-linux/fileshare/service"
+	"github.com/NordSecurity/nordvpn-linux/internal"
+	"github.com/NordSecurity/nordvpn-linux/test/category"
+	"github.com/NordSecurity/nordvpn-linux/test/mock/networker"
+	mapset "github.com/deckarep/golang-set"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRPCCities(t *testing.T) {
+	category.Set(t, category.Unit)
+	defer testsCleanup()
+
+	tests := []struct {
+		name       string
+		dm         *DataManager
+		cm         config.Manager
+		statusCode int64
+	}{
+		{
+			name:       "missing configuration file",
+			dm:         testNewDataManager(),
+			cm:         failingConfigManager{},
+			statusCode: internal.CodeConfigError,
+		},
+		{
+			name:       "app data is empty",
+			dm:         testNewDataManager(),
+			cm:         newMockConfigManager(),
+			statusCode: internal.CodeEmptyPayloadError,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			rpc := RPC{
+				ac:        &workingLoginChecker{},
+				cm:        test.cm,
+				dm:        test.dm,
+				fileshare: service.NoopFileshare{},
+				netw:      &networker.Mock{},
+				ncClient:  mockNC{},
+				publisher: &subs.Subject[string]{},
+				api:       mockApi{},
+			}
+			payload, _ := rpc.Cities(context.Background(), &pb.CitiesRequest{})
+
+			assert.Equal(t, test.statusCode, payload.Type)
+		})
+	}
+}
+
+func TestRPCCities_Successful(t *testing.T) {
+	category.Set(t, category.Unit)
+	defer testsCleanup()
+
+	dm := testNewDataManager()
+
+	cityNames := map[bool]map[config.Protocol]map[string]mapset.Set{
+		false: {
+			config.Protocol_UDP: {
+				"lt": mapset.NewSet("Vilnius"),
+			},
+			config.Protocol_TCP: {
+				"de": mapset.NewSet("Berlin"),
+			},
+		},
+	}
+	dm.SetAppData(nil, cityNames, nil)
+
+	cm := newMockConfigManager()
+	cm.c.AutoConnectData.Protocol = config.Protocol_UDP
+
+	rpc := RPC{
+		ac:        &workingLoginChecker{},
+		cm:        cm,
+		dm:        dm,
+		fileshare: service.NoopFileshare{},
+		netw:      &networker.Mock{},
+		ncClient:  mockNC{},
+		publisher: &subs.Subject[string]{},
+		api:       mockApi{},
+	}
+
+	request := &pb.CitiesRequest{}
+	request.Obfuscate = false
+	request.Country = "LT"
+
+	payload, _ := rpc.Cities(context.Background(), request)
+	assert.Equal(t, internal.CodeSuccess, payload.GetType())
+	assert.Equal(t, []string{"Vilnius"}, payload.GetData())
+}