1// Copyright (C) 2021 Storj Labs, Inc.
2// See LICENSE for copying information
3
4package geoip_test
5
6import (
7	"fmt"
8	"testing"
9
10	"github.com/stretchr/testify/assert"
11	"github.com/stretchr/testify/require"
12	"go.uber.org/zap"
13
14	"storj.io/common/storj/location"
15	"storj.io/common/testcontext"
16	"storj.io/storj/private/testplanet"
17	"storj.io/storj/satellite"
18	"storj.io/storj/storagenode"
19)
20
21func TestGeoIPMock(t *testing.T) {
22	testplanet.Run(t,
23		testplanet.Config{
24			SatelliteCount: 1, StorageNodeCount: 3, UplinkCount: 0,
25			Reconfigure: testplanet.Reconfigure{
26				Satellite: func(logger *zap.Logger, index int, config *satellite.Config) {
27					config.Overlay.GeoIP.MockCountries = []string{"US", "GB"}
28				},
29				StorageNode: func(index int, config *storagenode.Config) {
30					config.Server.Address = fmt.Sprintf("127.0.201.%d:0", index+1)
31				},
32			},
33		},
34		func(t *testing.T, ctx *testcontext.Context, planet *testplanet.Planet) {
35			// ensure storage nodes checked in with satellite
36			for _, node := range planet.StorageNodes {
37				node.Contact.Chore.TriggerWait(ctx)
38			}
39
40			// expected country codes per node index
41			countryCodes := map[int]location.CountryCode{
42				0: location.UnitedKingdom,
43				1: location.UnitedStates,
44				2: location.UnitedKingdom,
45			}
46
47			// check the country code for each storage nodes
48			for i, node := range planet.StorageNodes {
49				dossier, err := planet.Satellites[0].API.Overlay.DB.Get(ctx, node.ID())
50				require.NoError(t, err)
51				assert.Equal(t, countryCodes[i], dossier.CountryCode)
52			}
53		},
54	)
55}
56