-
Notifications
You must be signed in to change notification settings - Fork 390
/
ip2country_mock.go
41 lines (32 loc) · 1.14 KB
/
ip2country_mock.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
// Copyright (C) 2021 Storj Labs, Inc.
// See LICENSE for copying information
package geoip
import "storj.io/common/storj/location"
// MockIPToCountry provides a mock solution for looking up country codes in testplanet tests. This is done using the
// last byte of the ip address and mod'ing it into a country code.
type MockIPToCountry []location.CountryCode
// NewMockIPToCountry creates a mock IPToCountry based on predefined country list.
func NewMockIPToCountry(countries []string) MockIPToCountry {
result := MockIPToCountry{}
for _, country := range countries {
result = append(result, location.ToCountryCode(country))
}
return result
}
// Close does nothing for the MockIPToCountry.
func (m MockIPToCountry) Close() error {
return nil
}
// LookupISOCountryCode accepts an IP address.
func (m MockIPToCountry) LookupISOCountryCode(address string) (location.CountryCode, error) {
if len(m) == 0 {
return location.CountryCode(0), nil
}
ip, err := addressToIP(address)
if err != nil || ip == nil {
return location.CountryCode(0), err
}
lastBlock := int(ip[len(ip)-1])
return m[lastBlock%len(m)], nil
}
var _ IPToCountry = MockIPToCountry{}