Move code out of main and add tests for encoding

Author: oschwald

cmd/mmdbinspect/main.go

@@ -3,14 +3,11 @@
 package main
 
 import (
-	"encoding/json"
 	"flag"
 	"fmt"
 	"log"
 	"os"
 	"strings"
-
-	"github.com/goccy/go-yaml"
 )
 
 type arrayFlags []string
@@ -53,9 +50,9 @@ func main() {
 	flag.Parse()
 
 	// Any remaining arguments (not passed via flags) should be networks
-	network := flag.Args()
+	networks := flag.Args()
 
-	if len(network) == 0 {
+	if len(networks) == 0 {
 		fmt.Println("You must provide at least one network address")
 		usage()
 		os.Exit(1)
@@ -69,34 +66,17 @@ func main() {
 
 	w := os.Stdout
 
-	var encoder interface {
-		Encode(any) error
-	}
-	if *useJSONL {
-		enc := json.NewEncoder(w)
-		enc.SetEscapeHTML(false) // don't escape ampersands and angle brackets
-		encoder = enc
-	} else {
-		encoder = yaml.NewEncoder(w)
-	}
-
-	iterator := records(
-		network,
+	err := process(
+		w,
+		*useJSONL,
+		networks,
 		mmdb,
 		*includeAliasedNetworks,
 		*includeBuildTime,
 		*includeNetworksWithoutData,
 	)
-
-	for r, err := range iterator {
-		if err != nil {
-			log.Fatal(err)
-		}
-
-		err = encoder.Encode(r)
-		if err != nil {
-			log.Fatal(fmt.Errorf("encoding record: %w", err))
-		}
+	if err != nil {
+		log.Fatal(err)
 	}
 }
 

cmd/mmdbinspect/process.go

@@ -0,0 +1,50 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+
+	"github.com/goccy/go-yaml"
+)
+
+func process(
+	w io.Writer,
+	useJSONL bool,
+	networks, databases []string,
+	includeAliasedNetworks,
+	includeBuildTime,
+	includeNetworksWithoutData bool,
+) error {
+	var encoder interface {
+		Encode(any) error
+	}
+
+	if useJSONL {
+		enc := json.NewEncoder(w)
+		enc.SetEscapeHTML(false) // don't escape ampersands and angle brackets
+		encoder = enc
+	} else {
+		encoder = yaml.NewEncoder(w)
+	}
+
+	iterator := records(
+		networks,
+		databases,
+		includeAliasedNetworks,
+		includeBuildTime,
+		includeNetworksWithoutData,
+	)
+
+	for r, err := range iterator {
+		if err != nil {
+			return err
+		}
+
+		if err := encoder.Encode(r); err != nil {
+			return fmt.Errorf("encoding record: %w", err)
+		}
+	}
+
+	return nil
+}

cmd/mmdbinspect/process_test.go

@@ -0,0 +1,59 @@
+package main
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestProcess(t *testing.T) {
+	tests := []struct {
+		name      string
+		useJSONL  bool
+		networks  []string
+		databases []string
+		expected  string
+	}{
+		{
+			name:      "single record YAML",
+			useJSONL:  false,
+			networks:  []string{"81.2.69.142"},
+			databases: []string{CountryDBPath},
+			expected:  "database_path:.*\nrequested_lookup: 81.2.69.142\nnetwork: 81.2.69.142/31\nrecord:\n",
+		},
+		{
+			name:      "Single record JSONL",
+			useJSONL:  true,
+			networks:  []string{"81.2.69.142"},
+			databases: []string{CountryDBPath},
+			expected:  `{"database_path":"../../test/data/test-data/GeoIP2-Country-Test.mmdb","requested_lookup":"81.2.69.142","network":"81.2.69.142/31","record":{"continent":{"code":"EU","geoname_id":6255148,"names":{"de":"Europa","en":"Europe","es":"Europa","fr":"Europe","ja":"ヨーロッパ","pt-BR":"Europa","ru":"Европа","zh-CN":"欧洲"}},"country":{"geoname_id":2635167,"is_in_european_union":true,"iso_code":"GB","names":{"de":"Vereinigtes Königreich","en":"United Kingdom","es":"Reino Unido","fr":"Royaume-Uni","ja":"イギリス","pt-BR":"Reino Unido","ru":"Великобритания","zh-CN":"英国"}},"registered_country":{"geoname_id":6252001,"iso_code":"US","names":{"de":"USA","en":"United States","es":"Estados Unidos","fr":"États-Unis","ja":"アメリカ合衆国","pt-BR":"Estados Unidos","ru":"США","zh-CN":"美国"}}}}`,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			var buf bytes.Buffer
+
+			err := process(
+				&buf,
+				test.useJSONL,
+				test.networks,
+				test.databases,
+				false,
+				false,
+				false,
+			)
+
+			require.NoError(t, err)
+
+			if test.useJSONL {
+				assert.JSONEq(t, test.expected, buf.String())
+			} else {
+				assert.Regexp(t, test.expected, buf.String())
+			}
+		})
+	}
+}