package geolite import ( "archive/tar" "bytes" "compress/gzip" "io" "log/slog" "net/http" "os" "path/filepath" "sync/atomic" "testing" "github.com/stretchr/testify/require" testutils "github.com/pocket-id/pocket-id/backend/internal/utils/testing" ) // testDownloadURL is the URL the mock HTTP client serves the database from const testDownloadURL = "https://example.com/geolite/GeoLite2-City.tar.gz" // testDatabasePath is the sample database published by MaxMind, see testdata/README.md const testDatabasePath = "testdata/GeoLite2-City-Test.mmdb" // readTestDatabase returns the raw sample GeoLite2 City database func readTestDatabase(t *testing.T) []byte { t.Helper() data, err := os.ReadFile(testDatabasePath) require.NoError(t, err) return data } // testLogger returns a logger that discards everything, so tests don't spam the output func testLogger() *slog.Logger { return slog.New(slog.DiscardHandler) } // newServiceForTest returns a Service backed by a database file inside a temporary directory, along with the path of that file // When data is nil no database is written, so the service starts with nothing to look up against func newServiceForTest(t *testing.T, data []byte) (*Service, string) { t.Helper() dbPath := filepath.Join(t.TempDir(), "GeoLite2-City.mmdb") if data != nil { writeDatabaseFileForTest(t, dbPath, data) } return newServiceAtPathForTest(t, dbPath), dbPath } // newServiceAtPathForTest returns a Service backed by the database file at dbPath, loaded and unmapped when the test ends // Unmapping is required on Windows func newServiceAtPathForTest(t *testing.T, dbPath string) *Service { t.Helper() svc := newService(testLogger(), dbPath) t.Cleanup(svc.unload) err := svc.load(t.Context()) require.NoError(t, err) return svc } // writeDatabaseFileForTest puts a database at path the same way the refresher does: written elsewhere, then moved into place func writeDatabaseFileForTest(t *testing.T, path string, data []byte) { t.Helper() tmpPath := path + ".tmp" err := os.WriteFile(tmpPath, data, 0600) require.NoError(t, err) err = os.Rename(tmpPath, path) require.NoError(t, err) } // countingRoundTripper serves a fixed response for testDownloadURL and counts how many requests it has received type countingRoundTripper struct { body []byte statusCode int requests atomic.Int32 } func (rt *countingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { if req.URL.String() != testDownloadURL { return testutils.NewMockResponse(http.StatusNotFound, ""), nil } rt.requests.Add(1) statusCode := rt.statusCode if statusCode == 0 { statusCode = http.StatusOK } return &http.Response{ StatusCode: statusCode, Body: io.NopCloser(bytes.NewReader(rt.body)), Header: make(http.Header), ContentLength: int64(len(rt.body)), }, nil } // newDownloadClientForTest returns an HTTP client that serves body at testDownloadURL, along with the transport that counts the requests it receives func newDownloadClientForTest(body []byte) (*http.Client, *countingRoundTripper) { rt := &countingRoundTripper{body: body} return &http.Client{Transport: rt}, rt } // buildTarGzForTest returns a gzipped tarball holding the given files, mirroring the archive MaxMind publishes func buildTarGzForTest(t *testing.T, files map[string][]byte) []byte { t.Helper() buf := &bytes.Buffer{} gzw := gzip.NewWriter(buf) tw := tar.NewWriter(gzw) for name, content := range files { err := tw.WriteHeader(&tar.Header{ Name: name, Mode: 0o644, Size: int64(len(content)), Typeflag: tar.TypeReg, }) require.NoError(t, err) _, err = tw.Write(content) require.NoError(t, err) } require.NoError(t, tw.Close()) require.NoError(t, gzw.Close()) return buf.Bytes() }