beszel-ipv6/agent/handlers_test.go

//go:build testing

package agent

import (
	"testing"
	"time"

	"github.com/fxamacker/cbor/v2"
	"github.com/henrygd/beszel/agent/zfs"
	"github.com/henrygd/beszel/internal/common"
	"github.com/henrygd/beszel/internal/entities/smart"
	"github.com/henrygd/beszel/internal/entities/system"
	"github.com/stretchr/testify/assert"
)

// MockHandler for testing
type MockHandler struct {
	requiresVerification bool
	description          string
	handleFunc           func(ctx *HandlerContext) error
}

func TestNewAgentResponseSmartData(t *testing.T) {
	response := newAgentResponse(smart.SmartDataResponse{
		Data: map[string]smart.SmartData{
			"AAA": {SerialNumber: "AAA"},
		},
		Complete: true,
	}, nil)

	assert.Equal(t, "AAA", response.SmartData["AAA"].SerialNumber)
	assert.True(t, response.SmartComplete)
}

func TestGetDataHandlerReportsRequestTransport(t *testing.T) {
	cache := NewSystemDataCache()
	cached := &system.CombinedData{}
	cache.Set(cached, defaultDataCacheTimeMs)
	agent := &Agent{cache: cache}
	options, err := cbor.Marshal(common.DataRequestOptions{CacheTimeMs: defaultDataCacheTimeMs})
	assert.NoError(t, err)
	request := &common.HubRequest[cbor.RawMessage]{Action: common.GetData, Data: options}
	for _, transport := range []system.ConnectionType{system.ConnectionTypeSSH, system.ConnectionTypeWebSocket} {
		ctx := &HandlerContext{
			Agent:          agent,
			Request:        request,
			ConnectionType: transport,
			SendResponse: func(data any, _ *uint32) error {
				response := data.(*system.CombinedData)
				assert.Equal(t, transport, response.Info.ConnectionType)
				return nil
			},
		}
		assert.NoError(t, (&GetDataHandler{}).Handle(ctx))
		assert.Equal(t, system.ConnectionTypeNone, cached.Info.ConnectionType)
	}
}

func TestGetZfsDataHandlerForceRefresh(t *testing.T) {
	poolCalls := 0
	zm := &StoragePoolManager{detailInterval: time.Hour, backends: []*poolBackend{{name: "zfs"}}}
	zm.backends[0].poolStatsFn = func() ([]zfs.PoolStat, error) {
		poolCalls++
		return []zfs.PoolStat{{Name: "tank", Alloc: uint64(poolCalls)}}, nil
	}
	zm.backends[0].poolStatusesFn = func() ([]zfs.PoolStatus, error) { return nil, nil }
	zm.backends[0].datasetsFn = func() ([]zfs.Dataset, error) { return nil, nil }
	zm.GetDetail(false)

	requestData, err := cbor.Marshal(common.ZfsDataRequest{Force: true})
	assert.NoError(t, err)
	ctx := &HandlerContext{
		Agent: &Agent{storagePoolManager: zm},
		Request: &common.HubRequest[cbor.RawMessage]{
			Action: common.GetZfsData,
			Data:   requestData,
		},
		SendResponse: func(any, *uint32) error { return nil },
	}

	assert.NoError(t, (&GetZfsDataHandler{}).Handle(ctx))
	assert.Equal(t, 2, poolCalls)
}

func (m *MockHandler) Handle(ctx *HandlerContext) error {
	if m.handleFunc != nil {
		return m.handleFunc(ctx)
	}
	return nil
}

func (m *MockHandler) RequiresVerification() bool {
	return m.requiresVerification
}

// TestHandlerRegistry tests the handler registry functionality
func TestHandlerRegistry(t *testing.T) {
	t.Run("default registration", func(t *testing.T) {
		registry := NewHandlerRegistry()

		// Check default handlers are registered
		getDataHandler, exists := registry.GetHandler(common.GetData)
		assert.True(t, exists)
		assert.IsType(t, &GetDataHandler{}, getDataHandler)

		fingerprintHandler, exists := registry.GetHandler(common.CheckFingerprint)
		assert.True(t, exists)
		assert.IsType(t, &CheckFingerprintHandler{}, fingerprintHandler)
	})

	t.Run("custom handler registration", func(t *testing.T) {
		registry := NewHandlerRegistry()
		mockHandler := &MockHandler{
			requiresVerification: true,
			description:          "Test handler",
		}

		// Register a custom handler for a mock action
		const mockAction common.WebSocketAction = 99
		registry.Register(mockAction, mockHandler)

		// Verify registration
		handler, exists := registry.GetHandler(mockAction)
		assert.True(t, exists)
		assert.Equal(t, mockHandler, handler)
	})

	t.Run("unknown action", func(t *testing.T) {
		registry := NewHandlerRegistry()
		ctx := &HandlerContext{
			Request: &common.HubRequest[cbor.RawMessage]{
				Action: common.WebSocketAction(255), // Unknown action
			},
			HubVerified: true,
		}

		err := registry.Handle(ctx)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "unknown action: 255")
	})

	t.Run("verification required", func(t *testing.T) {
		registry := NewHandlerRegistry()
		ctx := &HandlerContext{
			Request: &common.HubRequest[cbor.RawMessage]{
				Action: common.GetData, // Requires verification
			},
			HubVerified: false, // Not verified
		}

		err := registry.Handle(ctx)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "hub not verified")
	})
}

// TestCheckFingerprintHandler tests the CheckFingerprint handler
func TestCheckFingerprintHandler(t *testing.T) {
	handler := &CheckFingerprintHandler{}

	t.Run("handle with invalid data", func(t *testing.T) {
		client := &WebSocketClient{}
		ctx := &HandlerContext{
			Client:      client,
			HubVerified: false,
			Request: &common.HubRequest[cbor.RawMessage]{
				Action: common.CheckFingerprint,
				Data:   cbor.RawMessage{}, // Empty/invalid data
			},
		}

		// Should fail to decode the fingerprint request
		err := handler.Handle(ctx)
		assert.Error(t, err)
	})
}