netbird/e2e/agentnetwork/management_test.go

//go:build e2e

package agentnetwork

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/netbirdio/netbird/e2e/harness"
	"github.com/netbirdio/netbird/shared/management/client/rest"
	"github.com/netbirdio/netbird/shared/management/http/api"
)

func ptr[T any](v T) *T { return &v }

// newProvider creates an OpenAI-catalog provider with a dummy key (these tests
// never call the upstream) and registers cleanup.
func newProvider(t *testing.T, ctx context.Context, name string) api.AgentNetworkProvider {
	t.Helper()
	prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
		Name:        name,
		ProviderId:  "openai_api",
		UpstreamUrl: "https://api.openai.com",
		ApiKey:      ptr("sk-dummy-e2e-key"),
	})
	require.NoError(t, err, "create provider %q", name)
	t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
	return prov
}

// requireClientError asserts err is a REST APIError with a 4xx status.
func requireClientError(t *testing.T, err error) {
	t.Helper()
	var apiErr *rest.APIError
	require.ErrorAs(t, err, &apiErr, "expected a REST APIError")
	assert.GreaterOrEqual(t, apiErr.StatusCode, 400, "expected a 4xx status")
	assert.Less(t, apiErr.StatusCode, 500, "expected a 4xx status")
}

// TestProviderLifecycle covers create → get → list → delete → 404 for every
// available real provider catalog (and a synthetic OpenAI provider when no
// provider keys are set), so each catalog's create and field round-trip is
// exercised. Create is offline — no upstream call — so this stays fast and
// burns no provider quota.
func TestProviderLifecycle(t *testing.T) {
	ctx := context.Background()

	cases := availableProviders()
	if len(cases) == 0 {
		cases = []providerCase{{
			name: "openai", catalogID: "openai_api", upstream: "https://api.openai.com",
			apiKey: "sk-dummy-e2e-key", model: "gpt-4o-mini", kind: harness.WireChat,
		}}
	}

	for _, pc := range cases {
		pc := pc
		t.Run(pc.name, func(t *testing.T) {
			req := providerRequest(pc)
			req.Name = "lc-" + pc.name
			prov, err := srv.CreateProvider(ctx, req)
			require.NoError(t, err, "create %s provider", pc.name)
			t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })

			assert.NotEmpty(t, prov.Id, "created provider must have an id")
			assert.Equal(t, pc.catalogID, prov.ProviderId, "catalog id must round-trip")
			assert.Equal(t, req.Name, prov.Name, "name must round-trip")
			assert.Equal(t, pc.upstream, prov.UpstreamUrl, "upstream must round-trip")

			got, err := srv.GetProvider(ctx, prov.Id)
			require.NoError(t, err, "get provider")
			assert.Equal(t, prov.Id, got.Id)

			list, err := srv.ListProviders(ctx)
			require.NoError(t, err, "list providers")
			var ids []string
			for _, p := range list {
				ids = append(ids, p.Id)
			}
			assert.Contains(t, ids, prov.Id, "created provider must appear in the list")

			require.NoError(t, srv.DeleteProvider(ctx, prov.Id), "delete provider")
			_, err = srv.GetProvider(ctx, prov.Id)
			requireClientError(t, err)
		})
	}
}

// TestProviderValidation exercises the create-time validation rules. These are
// uniform across catalogs (no per-provider required-field rules exist: a
// catalog-specific malformed value such as a Vertex key without the keyfile::
// prefix is accepted at create and only fails at the proxy), so the cases here
// are catalog-agnostic: missing API key, unknown catalog id, an invalid upstream
// URL, and a blank name.
func TestProviderValidation(t *testing.T) {
	ctx := context.Background()

	_, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
		Name:        "No Key",
		ProviderId:  "openai_api",
		UpstreamUrl: "https://api.openai.com",
	})
	requireClientError(t, err)

	_, err = srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
		Name:        "Unknown Catalog",
		ProviderId:  "totally_unknown_provider",
		UpstreamUrl: "https://example.com",
		ApiKey:      ptr("sk-dummy"),
	})
	requireClientError(t, err)

	_, err = srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
		Name:        "Bad Upstream",
		ProviderId:  "openai_api",
		UpstreamUrl: "not-a-url",
		ApiKey:      ptr("sk-dummy"),
	})
	requireClientError(t, err)

	_, err = srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
		Name:        "   ",
		ProviderId:  "openai_api",
		UpstreamUrl: "https://api.openai.com",
		ApiKey:      ptr("sk-dummy"),
	})
	requireClientError(t, err)
}

// TestSettingsRoundTrip flips the collection toggles and confirms the
// endpoint and proxy address stay immutable, then restores the original
// state. A second bootstrap attempt must be rejected as a conflict.
func TestSettingsRoundTrip(t *testing.T) {
	ctx := context.Background()

	// The package's TestMain bootstrapped the shared account's endpoint.
	before, err := srv.GetSettings(ctx)
	require.NoError(t, err, "get settings")
	require.NotEmpty(t, before.Endpoint, "settings must carry the bootstrapped endpoint")
	require.NotEmpty(t, before.ProxyAddress, "settings must carry the bootstrapped proxy address")

	require.NotNil(t, before.AccessLogRetentionDays, "bootstrapped settings must carry a retention")
	beforeRetention := *before.AccessLogRetentionDays

	flipped, err := srv.UpdateSettings(ctx, api.AgentNetworkSettingsRequest{
		Endpoint:               before.Endpoint,
		ProxyAddress:           before.ProxyAddress,
		EnableLogCollection:    !before.EnableLogCollection,
		EnablePromptCollection: !before.EnablePromptCollection,
		RedactPii:              !before.RedactPii,
		AccessLogRetentionDays: beforeRetention,
	})
	require.NoError(t, err, "update settings")
	assert.Equal(t, !before.EnableLogCollection, flipped.EnableLogCollection, "log collection toggle must flip")
	assert.Equal(t, !before.EnablePromptCollection, flipped.EnablePromptCollection, "prompt collection toggle must flip")
	require.NotNil(t, flipped.AccessLogRetentionDays)
	assert.Equal(t, beforeRetention, *flipped.AccessLogRetentionDays,
		"retention sent unchanged must round-trip, not reset to the zero value")
	assert.Equal(t, before.Endpoint, flipped.Endpoint, "endpoint must be immutable across updates")
	assert.Equal(t, before.ProxyAddress, flipped.ProxyAddress, "proxy address must be immutable across updates")

	// The account is already bootstrapped: a second bootstrap is a conflict,
	// whatever shape it asks for.
	_, err = srv.CreateSettings(ctx, api.AgentNetworkSettingsCreateRequest{
		Endpoint: ptr("attacker.cluster.invalid"),
	})
	requireClientError(t, err)

	// The identity fields ride along on the PUT as a required echo: a request
	// carrying a different endpoint is rejected without applying anything.
	_, err = srv.UpdateSettings(ctx, api.AgentNetworkSettingsRequest{
		Endpoint:               "other.cluster.invalid",
		ProxyAddress:           before.ProxyAddress,
		EnableLogCollection:    before.EnableLogCollection,
		EnablePromptCollection: before.EnablePromptCollection,
		RedactPii:              before.RedactPii,
		AccessLogRetentionDays: beforeRetention,
	})
	requireClientError(t, err)

	// Restore the original toggles.
	_, err = srv.UpdateSettings(ctx, api.AgentNetworkSettingsRequest{
		Endpoint:               before.Endpoint,
		ProxyAddress:           before.ProxyAddress,
		EnableLogCollection:    before.EnableLogCollection,
		EnablePromptCollection: before.EnablePromptCollection,
		RedactPii:              before.RedactPii,
		AccessLogRetentionDays: beforeRetention,
	})
	require.NoError(t, err, "restore settings")
}

// TestPolicyWindowFloor rejects an enabled limit below the 60s window floor and
// accepts one at the floor.
func TestPolicyWindowFloor(t *testing.T) {
	ctx := context.Background()

	grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-policy-grp"})
	require.NoError(t, err, "create source group")
	t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) })

	prov := newProvider(t, ctx, "Policy Provider")

	limits := func(window int64) *api.AgentNetworkPolicyLimits {
		return &api.AgentNetworkPolicyLimits{
			TokenLimit: api.AgentNetworkPolicyTokenLimit{
				Enabled:       true,
				GroupCap:      1000,
				UserCap:       1000,
				WindowSeconds: window,
			},
		}
	}

	_, err = srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
		Name:                   "e2e-below-floor",
		SourceGroups:           []string{grp.Id},
		DestinationProviderIds: []string{prov.Id},
		Limits:                 limits(30),
	})
	requireClientError(t, err)

	pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
		Name:                   "e2e-at-floor",
		SourceGroups:           []string{grp.Id},
		DestinationProviderIds: []string{prov.Id},
		Limits:                 limits(60),
	})
	require.NoError(t, err, "policy at the 60s floor must be accepted")
	assert.NotEmpty(t, pol.Id, "created policy must have an id")
	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
}

// TestConsumptionList confirms the read endpoint always returns an array, never
// a 404/500.
func TestConsumptionList(t *testing.T) {
	ctx := context.Background()

	rows, err := srv.ListConsumption(ctx)
	require.NoError(t, err, "consumption list must not error")
	assert.NotNil(t, rows, "consumption must be a JSON array (possibly empty)")
}