ip-info/internal/data/data_test.go

package data

import (
	"errors"
	"net"
	"os"
	"path/filepath"
	"sync"
	"sync/atomic"
	"testing"
)

// nested looks up a nested map value by successive string keys, returning "" if absent.
func nested(m map[string]any, keys ...string) string {
	var cur any = m
	for _, k := range keys {
		mm, ok := cur.(map[string]any)
		if !ok {
			return ""
		}
		cur = mm[k]
	}
	s, _ := cur.(string)
	return s
}

// parseIP is a test helper that panics on invalid input (used only with literals).
func parseIP(s string) net.IP {
	ip := net.ParseIP(s)
	if ip == nil {
		panic("invalid ip in test: " + s)
	}
	return ip
}

// The package's data helpers use relative paths ("data/..."), so tests must
// run with the working directory at the repository root. TestMain moves there
// once for the whole package.
func TestMain(m *testing.M) {
	if root, err := filepath.Abs("../.."); err == nil {
		_ = os.Chdir(root)
	}
	os.Exit(m.Run())
}

// freshDB returns a loaded IpDb, skipping the test when no mmdb is present.
func freshDB(t *testing.T) *IpDb {
	t.Helper()
	db := Ins()
	if !db.IsLoaded() {
		t.Skip("no mmdb present, skipping test that needs real data")
	}
	return db
}

func TestQueryValidIP(t *testing.T) {
	db := freshDB(t)

	data, netRec, err := db.Query("8.8.8.8")
	if err != nil {
		t.Fatalf("Query(8.8.8.8) unexpected error: %v", err)
	}
	if data == nil {
		t.Fatal("Query(8.8.8.8) returned nil data")
	}
	// 8.8.8.8 is Google (AS15169), US.
	if got := nested(data, "country", "iso_code"); got != "US" {
		t.Errorf("country.iso_code = %v, want US", got)
	}
	if netRec == nil {
		t.Error("expected a non-nil network record")
	} else if !netRec.Contains(parseIP("8.8.8.8")) {
		t.Errorf("network %s should contain 8.8.8.8", netRec)
	}
}

func TestQueryInvalidIP(t *testing.T) {
	db := freshDB(t)

	if _, _, err := db.Query("not-an-ip"); err == nil {
		t.Fatal("expected error for invalid IP, got nil")
	}
}

func TestQueryNotFound(t *testing.T) {
	db := freshDB(t)

	// 240.0.0.0/4 is reserved/unallocated and should not be in a public GeoIP DB.
	_, _, err := db.Query("240.0.0.1")
	if err == nil {
		t.Fatal("expected not-found error, got nil")
	}
	if !errors.Is(err, ErrNotFound) {
		t.Errorf("expected ErrNotFound sentinel, got %v", err)
	}
}

// TestQueryReloadConcurrent is a regression test for the data race between
// Query (reads d.r) and Reload (reassigns d.r / closes the old reader). It
// must pass under `go test -race`.
func TestQueryReloadConcurrent(t *testing.T) {
	db := freshDB(t)

	var reloadOK, queries int64
	var wg sync.WaitGroup

	for w := 0; w < 8; w++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := 0; i < 5000; i++ {
				_, _, _ = db.Query("8.8.8.8")
				atomic.AddInt64(&queries, 1)
			}
		}()
	}
	for r := 0; r < 4; r++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := 0; i < 30; i++ {
				if db.Reload() == nil {
					atomic.AddInt64(&reloadOK, 1)
				}
			}
		}()
	}
	wg.Wait()

	if reloadOK == 0 {
		t.Fatalf("expected reloads to succeed, got reloads=%d queries=%d", reloadOK, queries)
	}
}