fleet/server/service/debug_trace_sampler_test.go

package service

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/fleetdm/fleet/v4/server/fleet"
	mockds "github.com/fleetdm/fleet/v4/server/mock"
	"github.com/fleetdm/fleet/v4/server/platform/tracing"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
)

func discardLogger() *slog.Logger {
	return slog.New(slog.NewTextHandler(io.Discard, nil))
}

// adminAuthedRequest builds a request and primes the mockService so the debug auth middleware lets it through as a global
// admin.
func adminAuthedRequest(t *testing.T, method, target string, body string) (*mockService, *http.Request) {
	t.Helper()
	svc := &mockService{}
	svc.On("GetSessionByKey", mock.Anything, "fake_session_key").
		Return(&fleet.Session{UserID: 42, ID: 1}, nil)
	svc.On("UserUnauthorized", mock.Anything, uint(42)).
		Return(&fleet.User{ID: 42, GlobalRole: new(fleet.RoleAdmin)}, nil)

	var reqBody io.Reader
	if body != "" {
		reqBody = strings.NewReader(body)
	}
	req := httptest.NewRequest(method, target, reqBody)
	req.Header.Add("Authorization", "BEARER fake_session_key")
	return svc, req
}

func TestTraceSamplerHandler_GET(t *testing.T) {
	svc, req := adminAuthedRequest(t, http.MethodGet, "https://fleetdm.com/debug/trace_sampler", "")

	ds := new(mockds.Store)
	ds.GetTraceSamplerSettingsFunc = func(_ context.Context) (*tracing.Settings, error) {
		return &tracing.Settings{
			HighVolumeRatio: 0.001,
			StandardRatio:   0.02,
			ForceFull:       false,
		}, nil
	}

	handler := MakeDebugHandler(svc, testConfig, discardLogger(), nil, ds)
	res := httptest.NewRecorder()
	handler.ServeHTTP(res, req)

	require.Equal(t, http.StatusOK, res.Code)
	require.True(t, ds.GetTraceSamplerSettingsFuncInvoked)

	var got tracing.Settings
	require.NoError(t, json.Unmarshal(res.Body.Bytes(), &got))
	require.InDelta(t, 0.001, got.HighVolumeRatio, 1e-9)
	require.InDelta(t, 0.02, got.StandardRatio, 1e-9)
	require.False(t, got.ForceFull)
}

func TestTraceSamplerHandler_PATCH_PersistsChangesAndReturnsRow(t *testing.T) {
	svc, req := adminAuthedRequest(t, http.MethodPatch,
		"https://fleetdm.com/debug/trace_sampler",
		`{"force_full": true}`)

	ds := new(mockds.Store)
	ds.GetTraceSamplerSettingsFunc = func(_ context.Context) (*tracing.Settings, error) {
		return &tracing.Settings{
			HighVolumeRatio: 0.001,
			StandardRatio:   0.02,
			ForceFull:       false,
		}, nil
	}
	var saved *tracing.Settings
	ds.SetTraceSamplerSettingsFunc = func(_ context.Context, s *tracing.Settings) error {
		saved = s
		return nil
	}

	handler := MakeDebugHandler(svc, testConfig, discardLogger(), nil, ds)
	res := httptest.NewRecorder()
	handler.ServeHTTP(res, req)

	require.Equal(t, http.StatusOK, res.Code, "PATCH should return 200, body=%s", res.Body.String())
	require.True(t, ds.SetTraceSamplerSettingsFuncInvoked)
	require.NotNil(t, saved)
	require.True(t, saved.ForceFull, "force_full should now be true")
	require.InDelta(t, 0.001, saved.HighVolumeRatio, 1e-9, "other fields should be preserved")

	// Verify the response body matches what was saved. If we forgot to write the response, the test would still see 200 from
	// httptest's default but the body would be empty.
	var returned tracing.Settings
	require.NoError(t, json.Unmarshal(res.Body.Bytes(), &returned))
	require.True(t, returned.ForceFull)
	require.InDelta(t, saved.HighVolumeRatio, returned.HighVolumeRatio, 1e-9)
	require.InDelta(t, saved.StandardRatio, returned.StandardRatio, 1e-9)

	// PATCH response must NOT include updated_at. The handler reads the row before the write, so the pre-write timestamp
	// would be stale. Operators do a follow-up GET to see the post-write value.
	require.NotContains(t, res.Body.String(), "updated_at",
		"PATCH response must drop updated_at to avoid returning a stale timestamp")
}

func TestTraceSamplerHandler_PATCH_PartialUpdatePreservesOtherFields(t *testing.T) {
	// Locks in the docstring claim that "PATCH semantics mean only the provided fields are applied." Sending only
	// high_volume_ratio must leave standard_ratio and force_full at their prior values.
	svc, req := adminAuthedRequest(t, http.MethodPatch,
		"https://fleetdm.com/debug/trace_sampler",
		`{"high_volume_ratio": 0.5}`)

	ds := new(mockds.Store)
	ds.GetTraceSamplerSettingsFunc = func(_ context.Context) (*tracing.Settings, error) {
		return &tracing.Settings{
			HighVolumeRatio: 0.001,
			StandardRatio:   0.07,
			ForceFull:       true,
		}, nil
	}
	var saved *tracing.Settings
	ds.SetTraceSamplerSettingsFunc = func(_ context.Context, s *tracing.Settings) error {
		saved = s
		return nil
	}

	handler := MakeDebugHandler(svc, testConfig, discardLogger(), nil, ds)
	res := httptest.NewRecorder()
	handler.ServeHTTP(res, req)

	require.Equal(t, http.StatusOK, res.Code, "body=%s", res.Body.String())
	require.NotNil(t, saved)
	require.InDelta(t, 0.5, saved.HighVolumeRatio, 1e-9, "high_volume_ratio should be applied")
	require.InDelta(t, 0.07, saved.StandardRatio, 1e-9, "standard_ratio should be preserved from the prior row")
	require.True(t, saved.ForceFull, "force_full should be preserved from the prior row")
}

func TestTraceSamplerHandler_PATCH_ReadFailureReturns500(t *testing.T) {
	svc, req := adminAuthedRequest(t, http.MethodPatch,
		"https://fleetdm.com/debug/trace_sampler",
		`{"force_full": true}`)

	ds := new(mockds.Store)
	ds.GetTraceSamplerSettingsFunc = func(_ context.Context) (*tracing.Settings, error) {
		return nil, errors.New("db unavailable")
	}

	handler := MakeDebugHandler(svc, testConfig, discardLogger(), nil, ds)
	res := httptest.NewRecorder()
	handler.ServeHTTP(res, req)

	require.Equal(t, http.StatusInternalServerError, res.Code)
	require.False(t, ds.SetTraceSamplerSettingsFuncInvoked, "should not attempt to write when read fails")
}

func TestTraceSamplerHandler_PATCH_WriteFailureReturns500(t *testing.T) {
	svc, req := adminAuthedRequest(t, http.MethodPatch,
		"https://fleetdm.com/debug/trace_sampler",
		`{"force_full": true}`)

	ds := new(mockds.Store)
	ds.GetTraceSamplerSettingsFunc = func(_ context.Context) (*tracing.Settings, error) {
		return &tracing.Settings{HighVolumeRatio: 0.001, StandardRatio: 0.02}, nil
	}
	ds.SetTraceSamplerSettingsFunc = func(_ context.Context, _ *tracing.Settings) error {
		return errors.New("constraint violation")
	}

	handler := MakeDebugHandler(svc, testConfig, discardLogger(), nil, ds)
	res := httptest.NewRecorder()
	handler.ServeHTTP(res, req)

	require.Equal(t, http.StatusInternalServerError, res.Code)
	require.True(t, ds.SetTraceSamplerSettingsFuncInvoked)
}

func TestTraceSamplerHandler_PATCH_RejectsBadJSON(t *testing.T) {
	svc, req := adminAuthedRequest(t, http.MethodPatch,
		"https://fleetdm.com/debug/trace_sampler",
		`{"force_full":`) // malformed

	ds := new(mockds.Store)
	handler := MakeDebugHandler(svc, testConfig, discardLogger(), nil, ds)
	res := httptest.NewRecorder()
	handler.ServeHTTP(res, req)

	require.Equal(t, http.StatusBadRequest, res.Code)
	require.False(t, ds.SetTraceSamplerSettingsFuncInvoked)
}

func TestTraceSamplerHandler_PATCH_RejectsOutOfRangeRatio(t *testing.T) {
	cases := []struct {
		name string
		body string
	}{
		{"high above 1", `{"high_volume_ratio": 1.5}`},
		{"high below 0", `{"high_volume_ratio": -0.1}`},
		{"standard above 1", `{"standard_ratio": 2.0}`},
		{"standard below 0", `{"standard_ratio": -1.0}`},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			svc, req := adminAuthedRequest(t, http.MethodPatch,
				"https://fleetdm.com/debug/trace_sampler", c.body)
			ds := new(mockds.Store)
			handler := MakeDebugHandler(svc, testConfig, discardLogger(), nil, ds)
			res := httptest.NewRecorder()
			handler.ServeHTTP(res, req)

			require.Equal(t, http.StatusBadRequest, res.Code)
			require.Contains(t, res.Body.String(), "must be in [0, 1]")
			require.False(t, ds.SetTraceSamplerSettingsFuncInvoked)
		})
	}
}

func TestTraceSamplerHandler_PATCH_RequiresAtLeastOneField(t *testing.T) {
	svc, req := adminAuthedRequest(t, http.MethodPatch,
		"https://fleetdm.com/debug/trace_sampler", `{}`)
	ds := new(mockds.Store)
	handler := MakeDebugHandler(svc, testConfig, discardLogger(), nil, ds)
	res := httptest.NewRecorder()
	handler.ServeHTTP(res, req)

	require.Equal(t, http.StatusBadRequest, res.Code)
	require.Contains(t, res.Body.String(), "at least one")
}