github.com/jimpick/sp-kyc-checks@v0.0.0-20230201194251-fa84fca72da8/checks/geoip/check_geo_test.go (about)

     1  package geoip
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"io/ioutil"
     7  	"log"
     8  	"os"
     9  	"strconv"
    10  	"testing"
    11  
    12  	"github.com/stretchr/testify/assert"
    13  )
    14  
    15  type TestCase struct {
    16  	minerID     string
    17  	city        string
    18  	countryCode string
    19  	want        bool
    20  }
    21  
    22  func TestGeoMatchExists(t *testing.T) {
    23  	minerID := os.Getenv("MINER_ID")
    24  	city := os.Getenv("CITY")
    25  	countryCode := os.Getenv("COUNTRY_CODE")
    26  	extraArtifacts := os.Getenv("EXTRA_ARTIFACTS")
    27  
    28  	var currentEpoch int64
    29  
    30  	cases := make([]TestCase, 0)
    31  	if minerID == "" {
    32  		currentEpoch = 2055000 // To match JSON files
    33  		cases = append(
    34  			cases,
    35  			TestCase{
    36  				minerID:     "f02620",
    37  				city:        "Warsaw",
    38  				countryCode: "PL",
    39  				want:        true,
    40  			},
    41  			TestCase{
    42  				minerID:     "f02620",
    43  				city:        "Toronto",
    44  				countryCode: "CA",
    45  				want:        false,
    46  			},
    47  			TestCase{ // No IPs
    48  				minerID:     "f01095710",
    49  				city:        "Wuzheng",
    50  				countryCode: "CN",
    51  				want:        false,
    52  			},
    53  		)
    54  		if os.Getenv("MAXMIND_USER_ID") == "skip" {
    55  			log.Println("Warning: Skipping tests as MAXMIND_USER_ID set to 'skip'")
    56  		} else {
    57  			cases = append(
    58  				cases,
    59  				TestCase{ // No GeoLite2 city, but has GeoIP2 data
    60  					minerID:     "f01736668",
    61  					city:        "Omaha",
    62  					countryCode: "US",
    63  					want:        true,
    64  				},
    65  				TestCase{ // Has GeoIP2 data, but no city
    66  					minerID:     "f01873432",
    67  					city:        "Las Vegas",
    68  					countryCode: "US",
    69  					want:        true,
    70  				},
    71  				TestCase{ // Bad data for country
    72  					minerID:     "f01873432",
    73  					city:        "Las Vegas",
    74  					countryCode: "United States",
    75  					want:        true,
    76  				},
    77  				TestCase{ // Bad data for country
    78  					minerID:     "f01558688",
    79  					city:        "Montreal",
    80  					countryCode: "Canada",
    81  					want:        true,
    82  				},
    83  			)
    84  		}
    85  		if os.Getenv("GOOGLE_MAPS_API_KEY") == "skip" {
    86  			log.Println("Warning: Skipping tests as GOOGLE_MAP_API_KEY set to 'skip'")
    87  		} else {
    88  			cases = append(
    89  				cases,
    90  				TestCase{ // Distance match, 500km
    91  					minerID:     "f01558688",
    92  					city:        "Montreal",
    93  					countryCode: "CA",
    94  					want:        true,
    95  				},
    96  				TestCase{ // More than 500 km
    97  					minerID:     "f01558688",
    98  					city:        "Vancouver",
    99  					countryCode: "CA",
   100  					want:        false,
   101  				},
   102  				TestCase{ // China - City Name match
   103  					minerID:     "f01012",
   104  					city:        "Hangzhou",
   105  					countryCode: "CN",
   106  					want:        true,
   107  				},
   108  				TestCase{ // China - City Name match, lowercase country code
   109  					minerID:     "f01012",
   110  					city:        "Hangzhou",
   111  					countryCode: "cn",
   112  					want:        true,
   113  				},
   114  				TestCase{ // China - Distance match
   115  					minerID:     "f01012",
   116  					city:        "Jiaxing",
   117  					countryCode: "CN",
   118  					want:        true,
   119  				},
   120  				TestCase{ // China - GeoLite2, no Baidu
   121  					minerID:     "f01901765",
   122  					city:        "Hangzhou",
   123  					countryCode: "CN",
   124  					want:        true,
   125  				},
   126  			)
   127  		}
   128  	} else {
   129  		var err error
   130  		currentEpoch, err = strconv.ParseInt(os.Getenv("EPOCH"), 10, 64)
   131  		if currentEpoch == 0 || err != nil {
   132  			currentEpoch, err = GetCurrentEpoch(context.Background())
   133  			if err != nil {
   134  				log.Fatalf("Error getting current epoch: %v\n", err)
   135  			}
   136  		}
   137  
   138  		cases = append(cases, TestCase{minerID, city, countryCode, true})
   139  	}
   140  
   141  	geodata, err := LoadGeoData()
   142  	assert.Nil(t, err)
   143  
   144  	geocodeClient, err := getGeocodeClient()
   145  	assert.Nil(t, err)
   146  
   147  	for _, c := range cases {
   148  		ok, extra, err := GeoMatchExists(context.Background(), geodata, geocodeClient,
   149  			currentEpoch, c.minerID, c.city, c.countryCode)
   150  		assert.Nil(t, err)
   151  		assert.Equal(t, c.want, ok)
   152  		if extraArtifacts != "" {
   153  			extraJson, err := json.MarshalIndent(extra, "", "  ")
   154  			assert.Nil(t, err)
   155  			ioutil.WriteFile(extraArtifacts, extraJson, 0644)
   156  		}
   157  	}
   158  }