agentnetwork: fold in extra unit tests from review

Add coverage raised in review (originally generated in a separate PR):
- gRPC CheckLLMPolicyLimits handler: model forwarding, empty-model passthrough,
  model_blocked deny envelope, selector-error -> Internal, unconfigured ->
  Unimplemented.
- selector: guardrail-lookup error propagates, orphaned guardrail reference is
  unrestricted, model filter narrows candidates before cap scoring.
- proxy guardrail: enabled-but-empty allowlist and all-blank factory entries
  both deny every model for the provider.
This commit is contained in:
mlsmaycon
2026-07-26 13:13:17 +00:00
parent 692257a3eb
commit ef174b19e5
3 changed files with 249 additions and 0 deletions

View File

@@ -2,6 +2,7 @@ package agentnetwork
import (
"context"
"errors"
"testing"
"time"
@@ -246,3 +247,83 @@ func TestSelectPolicy_UnionAcrossPolicyGuardrails(t *testing.T) {
assert.True(t, res.Allow, "a model in any of the policy's allowlist guardrails must be permitted")
assert.Equal(t, "pol-A", res.SelectedPolicyID)
}
// TestSelectPolicy_GuardrailLookupErrorPropagates proves a store failure while
// resolving the candidate policies' guardrails surfaces as an error, not a
// silent allow/deny.
func TestSelectPolicy_GuardrailLookupErrorPropagates(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
expectPolicies(mockStore, "acc-1", policy)
mockStore.EXPECT().
GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), "acc-1").
Return(nil, errors.New("store unavailable"))
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "gpt-4o",
})
require.Error(t, err, "a guardrail-lookup failure must surface as an error")
assert.Nil(t, res)
}
// TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted proves a
// policy referencing a guardrail ID absent from the account's set (a stale
// reference) imposes no model restriction — same as no guardrail.
func TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-missing")
expectPolicies(mockStore, "acc-1", policy)
expectGuardrails(mockStore, "acc-1")
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "anything-goes",
})
require.NoError(t, err)
assert.True(t, res.Allow, "an orphaned guardrail reference must not restrict the model")
assert.Equal(t, "pol-A", res.SelectedPolicyID)
}
// TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter proves the model
// gate narrows candidates before cap scoring: the permitting policy is selected
// even though the blocked one has a larger, more attractive cap.
func TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter(t *testing.T) {
ctrl := gomock.NewController(t)
mgr, mockStore := newSelectorMgr(t, ctrl)
polBig := guardedPolicy("pol-big", "acc-1", []string{"grp-eng"}, "prov-1", "g-restrict")
polBig.Limits = types.PolicyLimits{
TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 1_000_000, WindowSeconds: 3600},
}
polSmall := guardedPolicy("pol-small", "acc-1", []string{"grp-eng"}, "prov-1", "g-permit")
polSmall.Limits = types.PolicyLimits{
TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 100, WindowSeconds: 3600},
}
expectPolicies(mockStore, "acc-1", polBig, polSmall)
expectGuardrails(mockStore, "acc-1",
allowlistGuardrail("g-restrict", "acc-1", "gpt-4o"),
allowlistGuardrail("g-permit", "acc-1", "claude-opus-4"),
)
expectConsumptionBatch(mockStore, nil)
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
AccountID: "acc-1",
GroupIDs: []string{"grp-eng"},
ProviderID: "prov-1",
Model: "claude-opus-4", // only pol-small's guardrail permits this
})
require.NoError(t, err)
assert.True(t, res.Allow)
assert.Equal(t, "pol-small", res.SelectedPolicyID,
"the model filter must exclude pol-big before cap scoring")
}

View File

@@ -0,0 +1,138 @@
package grpc
import (
"context"
"errors"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc/codes"
grpcstatus "google.golang.org/grpc/status"
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
"github.com/netbirdio/netbird/shared/management/proto"
)
// fakeAgentNetworkLimits records the PolicySelectionInput it was invoked with
// and returns a pre-programmed result, so tests can assert what the handler
// forwards to the selector.
type fakeAgentNetworkLimits struct {
gotInput agentnetwork.PolicySelectionInput
result *agentnetwork.PolicySelectionResult
err error
}
func (f *fakeAgentNetworkLimits) SelectPolicyForRequest(_ context.Context, in agentnetwork.PolicySelectionInput) (*agentnetwork.PolicySelectionResult, error) {
f.gotInput = in
if f.err != nil {
return nil, f.err
}
return f.result, nil
}
func (f *fakeAgentNetworkLimits) RecordUsage(_ context.Context, _ agentnetwork.RecordUsageInput) error {
return nil
}
// TestCheckLLMPolicyLimits_ForwardsModelToSelector proves the wiring added here:
// the model the proxy extracted must reach the selector's Model unchanged,
// alongside the account/user/group/provider fields.
func TestCheckLLMPolicyLimits_ForwardsModelToSelector(t *testing.T) {
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{Allow: true, SelectedPolicyID: "pol-1"}}
s := &ProxyServiceServer{}
s.SetAgentNetworkLimitsService(fake)
req := &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
UserId: "user-1",
GroupIds: []string{"grp-a", "grp-b"},
ProviderId: "prov-1",
Model: "claude-opus-4",
}
resp, err := s.CheckLLMPolicyLimits(context.Background(), req)
require.NoError(t, err)
require.NotNil(t, resp)
assert.Equal(t, "acc-1", fake.gotInput.AccountID)
assert.Equal(t, "user-1", fake.gotInput.UserID)
assert.Equal(t, []string{"grp-a", "grp-b"}, fake.gotInput.GroupIDs)
assert.Equal(t, "prov-1", fake.gotInput.ProviderID)
assert.Equal(t, "claude-opus-4", fake.gotInput.Model,
"the request's model must be forwarded to the selector")
}
// TestCheckLLMPolicyLimits_EmptyModelForwardedAsEmpty proves an undetermined
// model (empty string) is forwarded as-is; the selector decides how to treat it.
func TestCheckLLMPolicyLimits_EmptyModelForwardedAsEmpty(t *testing.T) {
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{Allow: true}}
s := &ProxyServiceServer{}
s.SetAgentNetworkLimitsService(fake)
req := &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
ProviderId: "prov-1",
}
_, err := s.CheckLLMPolicyLimits(context.Background(), req)
require.NoError(t, err)
assert.Equal(t, "", fake.gotInput.Model, "an absent model must be forwarded as empty, not fabricated")
}
// TestCheckLLMPolicyLimits_DenyResponseCarriesModelBlockedCode proves the deny
// envelope surfaces the model-allowlist deny code + reason through the response.
func TestCheckLLMPolicyLimits_DenyResponseCarriesModelBlockedCode(t *testing.T) {
fake := &fakeAgentNetworkLimits{result: &agentnetwork.PolicySelectionResult{
Allow: false,
DenyCode: "llm_policy.model_blocked",
DenyReason: `model "claude-opus-4" is not permitted by any applicable policy allowlist`,
}}
s := &ProxyServiceServer{}
s.SetAgentNetworkLimitsService(fake)
resp, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
ProviderId: "prov-1",
Model: "claude-opus-4",
})
require.NoError(t, err)
require.NotNil(t, resp)
assert.Equal(t, "deny", resp.Decision)
assert.Equal(t, "llm_policy.model_blocked", resp.DenyCode)
assert.NotEmpty(t, resp.DenyReason)
assert.Empty(t, resp.SelectedPolicyId, "a denied request must carry no selected policy")
}
// TestCheckLLMPolicyLimits_SelectorErrorSurfacesAsInternal proves a selector
// failure surfaces as an Internal gRPC error rather than a silent allow.
func TestCheckLLMPolicyLimits_SelectorErrorSurfacesAsInternal(t *testing.T) {
fake := &fakeAgentNetworkLimits{err: errors.New("boom")}
s := &ProxyServiceServer{}
s.SetAgentNetworkLimitsService(fake)
_, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
ProviderId: "prov-1",
Model: "gpt-4o",
})
require.Error(t, err)
st, ok := grpcstatus.FromError(err)
require.True(t, ok)
assert.Equal(t, codes.Internal, st.Code(), "selector errors must never fail open on the hot path")
}
// TestCheckLLMPolicyLimits_UnconfiguredServiceReturnsUnimplemented locks the
// fallback: with no limits service wired the RPC returns Unimplemented.
func TestCheckLLMPolicyLimits_UnconfiguredServiceReturnsUnimplemented(t *testing.T) {
s := &ProxyServiceServer{}
_, err := s.CheckLLMPolicyLimits(context.Background(), &proto.CheckLLMPolicyLimitsRequest{
AccountId: "acc-1",
ProviderId: "prov-1",
})
require.Error(t, err)
st, ok := grpcstatus.FromError(err)
require.True(t, ok)
assert.Equal(t, codes.Unimplemented, st.Code())
}

View File

@@ -207,6 +207,36 @@ func TestRestrictionsButNoResolvedProviderFailsClosed(t *testing.T) {
assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
}
func TestEnabledButEmptyAllowlistDeniesEveryModel(t *testing.T) {
// An allowlist-enabled provider with zero models is distinct from an
// unrestricted (absent) provider: it must deny every model.
mw := New(providerCfg())
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
))
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionDeny, out.Decision, "an enabled-but-empty allowlist must deny every model")
require.NotNil(t, out.DenyReason)
assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "deny code must be model_blocked, not model_unknown")
}
func TestFactoryAllEmptyEntriesDenyEveryModel(t *testing.T) {
// All the provider's entries are blank; they collapse to a non-nil empty
// list (deny everything for that provider), not "no restriction".
raw := []byte(`{"provider_allowlists":{"prov-1":[""," "]}}`)
mw, err := Factory{}.New(raw)
require.NoError(t, err)
out, err := mw.Invoke(context.Background(), newInputProvider(testProvider,
middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
))
require.NoError(t, err)
require.NotNil(t, out)
assert.Equal(t, middleware.DecisionDeny, out.Decision, "all-blank allowlist entries must still restrict the provider")
require.NotNil(t, out.DenyReason)
assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "deny code must be model_blocked")
}
func TestPromptCaptureDisabledEmitsNoPrompt(t *testing.T) {
mw := New(Config{})
out, err := mw.Invoke(context.Background(), newInput(