add tests.

.gitignore

@@ -12,4 +12,4 @@
 *.out
 .idea/
 .local_dist/
-ip-ranges.json
+coverage.txt

go.mod

@@ -4,6 +4,8 @@ require (
 	github.com/cheekybits/is v0.0.0-20150225183255-68e9c0620927 // indirect
 	github.com/kr/pretty v0.1.0 // indirect
 	github.com/matryer/try v0.0.0-20161228173917-9ac251b645a2
+	github.com/stretchr/testify v1.3.0
+	github.com/urfave/cli v1.20.0
 	gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect
 	gopkg.in/urfave/cli.v1 v1.20.0
 	gopkg.in/yaml.v2 v2.2.2

go.sum

@@ -1,12 +1,21 @@
 github.com/cheekybits/is v0.0.0-20150225183255-68e9c0620927 h1:SKI1/fuSdodxmNNyVBR8d7X/HuLnRpvvFO0AgyQk764=
 github.com/cheekybits/is v0.0.0-20150225183255-68e9c0620927/go.mod h1:h/aW8ynjgkuj+NQRlZcDbAbM1ORAbXjXX77sX7T289U=
+github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/matryer/try v0.0.0-20161228173917-9ac251b645a2 h1:JAEbJn3j/FrhdWA9jW8B5ajsLIjeuEHLi8xE4fk997o=
 github.com/matryer/try v0.0.0-20161228173917-9ac251b645a2/go.mod h1:0KeJpeMD6o+O4hW7qJOT7vyQPKrWmj26uf5wMc/IiIs=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q=
+github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
+github.com/urfave/cli v1.20.0 h1:fDqGv3UG/4jbVl/QkFwEdddtEDjh/5Ov6X+0B/3bPaw=
+github.com/urfave/cli v1.20.0/go.mod h1:70zkFmudgCuE/ngEzBv17Jvp/497gISqfk5gWijbERA=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=

loader.go

@@ -12,56 +12,57 @@ import (
 	"github.com/matryer/try"
 )
 
-func loadRanges(filePath string) (ranges IPRangeDoc, err error) {
-	if filePath != "" {
-		var fileContent []byte
-		fileContent, err = ioutil.ReadFile(filePath)
-		if err != nil {
-			return
-		}
-		err = json.Unmarshal(fileContent, &ranges)
-		if err != nil {
-			return
-		}
-	} else {
-		var request *http.Request
-		request, err = http.NewRequest(http.MethodGet, ipURL, nil)
-		if err != nil {
-			return
-		}
+func loadFromFile(filePath string) (ranges IPRangeDoc, err error) {
+	var fileContent []byte
+	fileContent, err = ioutil.ReadFile(filePath)
+	if err != nil {
+		return
+	}
+	err = json.Unmarshal(fileContent, &ranges)
+	return
+}
 
-		client := &http.Client{
-			Transport: &http.Transport{
-				MaxIdleConnsPerHost: 3,
-				DisableKeepAlives:   false,
-			},
-			Timeout: time.Duration(6) * time.Second,
-		}
-		var resp *http.Response
+func loadFromURL() (ranges IPRangeDoc, err error) {
+	var request *http.Request
+	request, err = http.NewRequest(http.MethodGet, ipURL, nil)
+	if err != nil {
+		return
+	}
+
+	client := &http.Client{
+		Transport: &http.Transport{
+			MaxIdleConnsPerHost: 3,
+			DisableKeepAlives:   false,
+		},
+		Timeout: time.Duration(6) * time.Second,
+	}
+	var resp *http.Response
 
-		err = try.Do(func(attempt int) (bool, error) {
-			var rErr error
-			resp, rErr = client.Do(request)
-			return attempt < 3, rErr
-		})
+	err = try.Do(func(attempt int) (bool, error) {
+		var rErr error
+		resp, rErr = client.Do(request)
+		return attempt < 3, rErr
+	})
 
-		if err != nil {
-			return
-		}
-		defer resp.Body.Close()
+	if err != nil {
+		return
+	}
+	defer resp.Body.Close()
 
-		// process response body
-		var syncRespBodyBytes []byte
-		syncRespBodyBytes, err = getResponseBody(resp)
-		if err != nil {
-			return
-		}
-		err = json.Unmarshal(syncRespBodyBytes, &ranges)
-		if err != nil {
-			return
-		}
+	var syncRespBodyBytes []byte
+	syncRespBodyBytes, err = getResponseBody(resp)
+	if err != nil {
+		return
+	}
+	err = json.Unmarshal(syncRespBodyBytes, &ranges)
+	return
+}
+
+func loadRanges(filePath string) (ranges IPRangeDoc, err error) {
+	if filePath != "" {
+		return loadFromFile(filePath)
 	}
-	return ranges, err
+	return loadFromURL()
 
 }
 

loader_test.go

@@ -0,0 +1,25 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLoadJSONFile(t *testing.T) {
+	ranges, err := loadFromFile("testdata/ip-ranges.json")
+	assert.NoError(t, err)
+	assert.NotEqual(t, 0, len(ranges.Prefixes))
+	assert.NotEqual(t, 0, len(ranges.IPv6Prefixes))
+	assert.NotEmpty(t, ranges.SyncToken)
+	assert.NotEmpty(t, ranges.CreateDate)
+}
+
+func TestLoadURL(t *testing.T) {
+	ranges, err := loadFromURL()
+	assert.NoError(t, err)
+	assert.NotEqual(t, 0, len(ranges.Prefixes))
+	assert.NotEqual(t, 0, len(ranges.IPv6Prefixes))
+	assert.NotEmpty(t, ranges.SyncToken)
+	assert.NotEmpty(t, ranges.CreateDate)
+}

lookup_test.go

@@ -0,0 +1,53 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLookupIP(t *testing.T) {
+	prefixes := []Prefix{
+		{
+			IPPrefix: "192.168.0.0/24", Region: "eu-west-1", Service: "EC2",
+		},
+		{
+			IPPrefix: "192.168.1.0/24", Region: "eu-west-2", Service: "S3",
+		},
+		{
+			IPPrefix: "192.168.2.0/24", Region: "eu-west-3", Service: "CLOUDFRONT",
+		},
+	}
+	IPv6Prefixes := []IPv6Prefix{
+		{
+			IPv6Prefix: "2001:800::/21", Region: "us-east-2", Service: "GLOBALACCELERATOR",
+		},
+		{
+			IPv6Prefix: "2600:1fa0:4000::/40", Region: "us-east-1", Service: "CODEBUILD",
+		},
+	}
+	doc := IPRangeDoc{
+		Prefixes:     prefixes,
+		IPv6Prefixes: IPv6Prefixes,
+	}
+
+	lpi := lookupPrefixInput{
+		doc: doc,
+		ip:  "192.168.0.100",
+	}
+	lpo, err := lookupPrefix(lpi)
+	assert.NoError(t, err)
+	assert.Equal(t, lpo.doc.Prefixes[0].IPPrefix, "192.168.0.0/24")
+	assert.Equal(t, lpo.doc.Prefixes[0].Region, "eu-west-1")
+	assert.Equal(t, lpo.doc.Prefixes[0].Service, "EC2")
+
+	l6pi := lookupPrefixInput{
+		doc: doc,
+		ip:  "2600:1fa0:4060:c11:34da:e8e9::",
+	}
+	lpo, err = lookupPrefix(l6pi)
+	assert.NoError(t, err)
+	assert.Equal(t, lpo.doc.IPv6Prefixes[0].IPv6Prefix, "2600:1fa0:4000::/40")
+	assert.Equal(t, lpo.doc.IPv6Prefixes[0].Region, "us-east-1")
+	assert.Equal(t, lpo.doc.IPv6Prefixes[0].Service, "CODEBUILD")
+}

main.go

@@ -2,12 +2,13 @@ package main
 
 import (
 	"fmt"
-	"gopkg.in/urfave/cli.v1"
 	"net"
 	"os"
 	"sort"
 	"strings"
 	"time"
+
+	"gopkg.in/urfave/cli.v1"
 )
 
 const (