OpenFlare/backend/openflare/plugins/server/domain/fleet/relay/logics_test.go

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

package relay

import (
	"context"
	"encoding/json"
	"testing"
	"time"

	"Wavelet/openflare/plugins/server/kernel/repository"
	"Wavelet/openflare/plugins/server/kernel/testhelper"

	"Wavelet/openflare/plugins/server/domain/fleet/agent"
	"Wavelet/openflare/plugins/server/kernel/model"

	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

func setupRelayTestDB(t *testing.T) func() {
	t.Helper()

	sqliteDB, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{
		DisableForeignKeyConstraintWhenMigrating: true,
	})
	require.NoError(t, err)
	require.NoError(t, sqliteDB.AutoMigrate(
		&model.OpenFlareNode{},
		&model.SystemConfig{},
		&model.OpenFlareNodeSystemProfile{},
		&model.OpenFlareMetricSnapshot{},
		&model.OpenFlareHealthEvent{},
		&model.OpenFlareNodeObservationFrps{},
	))

	repository.SetDBForTest(sqliteDB)
	agent.ResetAuthCacheForTest()
	testhelper.SetupLogStoresForTest(t)

	return func() {
		repository.SetDBForTest(nil)
		agent.ResetAuthCacheForTest()
	}
}

func TestHeartbeatPayloadBindingAndFrpsObservationInsert(t *testing.T) {
	cleanup := setupRelayTestDB(t)
	defer cleanup()

	ctx := context.Background()
	now := time.Now().UTC().Truncate(time.Second)

	node := &model.OpenFlareNode{
		NodeID:      "node-relay-observe",
		Name:        "relay-1",
		AccessToken: "relay-token",
		Status:      "pending",
		NodeType:    "tunnel_relay",
		RelayStatus: "unknown",
	}
	require.NoError(t, repository.DB(ctx).Create(node).Error)

	proxies := []ProxyStat{
		{
			Name:          "proxy-a",
			Type:          "http",
			Status:        "online",
			ClientVersion: "0.61.0",
			ClientAddr:    "10.0.0.2:12345",
		},
	}
	_, err := Heartbeat(ctx, node, HeartbeatPayload{
		Version:         "v0.1.0",
		ExtVersion:      "0.61.0",
		RelayStatus:     "healthy",
		FrpsConnCount:   7,
		FrpsProxyCount:  3,
		FrpsClientCount: 2,
		FrpsProxies:     proxies,
		Name:            "relay-runtime",
		IP:              "203.0.113.9",
		Profile: &agent.NodeSystemProfile{
			Hostname:       "relay-runtime",
			OSName:         "Ubuntu",
			OSVersion:      "24.04",
			Architecture:   "amd64",
			CPUCores:       4,
			ReportedAtUnix: now.Unix(),
		},
		Snapshot: &agent.NodeMetricSnapshot{
			CapturedAtUnix:  now.Unix(),
			CPUUsagePercent: 12.5,
			DiskReadBytes:   100,
			DiskWriteBytes:  200,
		},
		HealthEvents: []agent.NodeHealthEvent{},
	})
	require.NoError(t, err)

	var stored model.OpenFlareNode
	require.NoError(t, repository.DB(ctx).Where("node_id = ?", node.NodeID).First(&stored).Error)
	assert.Equal(t, "online", stored.Status)
	assert.Equal(t, "healthy", stored.RelayStatus)
	assert.Equal(t, "203.0.113.9", stored.IP)
	assert.Equal(t, "v0.1.0", stored.Version)
	assert.Equal(t, "0.61.0", stored.ExtVersion)

	profile, err := repository.GetOpenFlareNodeSystemProfile(ctx, node.NodeID)
	require.NoError(t, err)
	assert.Equal(t, "relay-runtime", profile.Hostname)
	assert.Equal(t, "Ubuntu", profile.OSName)

	snapshots, err := repository.ListOpenFlareMetricSnapshotsSince(ctx, node.NodeID, now.Add(-time.Minute), 10)
	require.NoError(t, err)
	require.Len(t, snapshots, 1)
	assert.InDelta(t, 12.5, snapshots[0].CPUUsagePercent, 1e-9)

	frpsObs, err := repository.ListOpenFlareNodeObservationFrps(ctx, node.NodeID, time.Time{}, 1)
	require.NoError(t, err)
	require.Len(t, frpsObs, 1)
	assert.Equal(t, 7, frpsObs[0].FrpsConnections)
	assert.Equal(t, 3, frpsObs[0].FrpsProxyCount)
	assert.Equal(t, 2, frpsObs[0].FrpsClientCount)

	var decoded []ProxyStat
	require.NoError(t, json.Unmarshal([]byte(frpsObs[0].FrpsProxies), &decoded))
	require.Len(t, decoded, 1)
	assert.Equal(t, "proxy-a", decoded[0].Name)
	assert.Equal(t, "online", decoded[0].Status)
}

func TestHeartbeatRelayReconcilesFrpsUnhealthyEvent(t *testing.T) {
	cleanup := setupRelayTestDB(t)
	defer cleanup()

	ctx := context.Background()
	node := &model.OpenFlareNode{
		NodeID:      "node-relay-unhealthy",
		Name:        "relay-unhealthy",
		AccessToken: "relay-token-unhealthy",
		Status:      "pending",
		NodeType:    "tunnel_relay",
		RelayStatus: "healthy",
	}
	require.NoError(t, repository.DB(ctx).Create(node).Error)

	_, err := Heartbeat(ctx, node, HeartbeatPayload{
		Version:     "v0.1.0",
		ExtVersion:  "0.61.0",
		RelayStatus: "unhealthy",
	})
	require.NoError(t, err)

	events, err := repository.ListOpenFlareHealthEvents(ctx, node.NodeID, true, 10)
	require.NoError(t, err)
	require.Len(t, events, 1)
	assert.Equal(t, relayFrpsUnhealthyEventType, events[0].EventType)
	assert.Equal(t, "active", events[0].Status)

	_, err = Heartbeat(ctx, node, HeartbeatPayload{
		Version:     "v0.1.0",
		ExtVersion:  "0.61.0",
		RelayStatus: "healthy",
	})
	require.NoError(t, err)

	events, err = repository.ListOpenFlareHealthEvents(ctx, node.NodeID, false, 10)
	require.NoError(t, err)
	require.Len(t, events, 1)
	assert.Equal(t, "resolved", events[0].Status)
}