OpenFlare/internal/apps/openflare/agent/helpers_test.go

// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0

package agent

import (
	"context"
	"net"
	"testing"

	ofgeoip "github.com/Rain-kl/Wavelet/internal/apps/openflare/geoip"
	"github.com/Rain-kl/Wavelet/internal/model"
	pkggeoip "github.com/Rain-kl/Wavelet/pkg/geoip"
)

type fakeGeoIPProvider struct {
	info *pkggeoip.GeoInfo
}

func (f *fakeGeoIPProvider) Name() string { return "fake-geoip" }

func (f *fakeGeoIPProvider) GetGeoInfo(ip net.IP) (*pkggeoip.GeoInfo, error) {
	return f.info, nil
}

func (f *fakeGeoIPProvider) UpdateDatabase() error { return nil }

func (f *fakeGeoIPProvider) Close() error { return nil }

func withFakeGeoIPProvider(t *testing.T, info *pkggeoip.GeoInfo) {
	t.Helper()
	previous := pkggeoip.CurrentProvider
	pkggeoip.CurrentProvider = &fakeGeoIPProvider{info: info}
	t.Cleanup(func() {
		pkggeoip.CurrentProvider = previous
	})
}

func geoipFloat(value float64) *float64 {
	return &value
}

func TestApplyGeoInfoFromIP(t *testing.T) {
	latitude := 31.2304
	longitude := 121.4737
	withFakeGeoIPProvider(t, &pkggeoip.GeoInfo{
		Name:      "Shanghai",
		Latitude:  geoipFloat(latitude),
		Longitude: geoipFloat(longitude),
	})

	node := &model.OpenFlareNode{IP: "203.0.113.10"}
	ofgeoip.ApplyNodeGeoFromIP(context.Background(), node, node.IP)

	if node.GeoName != "Shanghai" {
		t.Fatalf("expected geo_name Shanghai, got %q", node.GeoName)
	}
	if node.GeoLatitude == nil || *node.GeoLatitude != latitude {
		t.Fatalf("unexpected geo_latitude: %+v", node.GeoLatitude)
	}
	if node.GeoLongitude == nil || *node.GeoLongitude != longitude {
		t.Fatalf("unexpected geo_longitude: %+v", node.GeoLongitude)
	}
}

func TestApplyGeoInfoFromIPSkipsInvalidIP(t *testing.T) {
	withFakeGeoIPProvider(t, &pkggeoip.GeoInfo{Name: "Should Not Apply"})

	node := &model.OpenFlareNode{
		IP:           "203.0.113.10",
		GeoName:      "Existing",
		GeoLatitude:  geoipFloat(1),
		GeoLongitude: geoipFloat(2),
	}
	ofgeoip.ApplyNodeGeoFromIP(context.Background(), node, "not-an-ip")

	if node.GeoName != "" || node.GeoLatitude != nil || node.GeoLongitude != nil {
		t.Fatalf("expected geo fields to be cleared on invalid IP, got %+v", node)
	}
}

func TestApplyNodeRuntimeRespectsGeoManualOverride(t *testing.T) {
	withFakeGeoIPProvider(t, &pkggeoip.GeoInfo{
		Name:      "Shanghai",
		Latitude:  geoipFloat(31.2304),
		Longitude: geoipFloat(121.4737),
	})

	node := &model.OpenFlareNode{
		GeoManualOverride: true,
		GeoName:           "Manual",
		GeoLatitude:       geoipFloat(10),
		GeoLongitude:      geoipFloat(20),
	}
	applyNodeRuntime(context.Background(), node, NodePayload{
		IP:      "203.0.113.10",
		Version: "1.0.0",
	}, true)

	if node.GeoName != "Manual" {
		t.Fatalf("expected manual geo_name to be preserved, got %q", node.GeoName)
	}
	if node.GeoLatitude == nil || *node.GeoLatitude != 10 {
		t.Fatalf("expected manual geo_latitude to be preserved, got %+v", node.GeoLatitude)
	}
}

func TestCollectHeartbeatChangesTracksGeoFields(t *testing.T) {
	before := &model.OpenFlareNode{
		IP:      "10.0.0.1",
		GeoName: "Old Region",
	}
	after := &model.OpenFlareNode{
		IP:           "203.0.113.10",
		GeoName:      "New Region",
		GeoLatitude:  geoipFloat(31.2304),
		GeoLongitude: geoipFloat(121.4737),
	}

	changes := collectHeartbeatChanges(before, after)
	if changes["ip"] != after.IP {
		t.Fatalf("expected ip change, got %+v", changes)
	}
	if changes["geo_name"] != after.GeoName {
		t.Fatalf("expected geo_name change, got %+v", changes)
	}
	if changes["geo_latitude"] != after.GeoLatitude {
		t.Fatalf("expected geo_latitude change, got %+v", changes)
	}
	if changes["geo_longitude"] != after.GeoLongitude {
		t.Fatalf("expected geo_longitude change, got %+v", changes)
	}
}