50 lines
2.3 KiB
Go
50 lines
2.3 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"git.cin.su/m/Assistant/contracts/capability"
|
|
"git.cin.su/m/Assistant/contracts/cognitive"
|
|
)
|
|
|
|
type admissionTestProvider struct{ calls int }
|
|
|
|
func (p *admissionTestProvider) Advise(context.Context, cognitive.AdvisoryRequest) (cognitive.AdvisoryResponse, error) {
|
|
p.calls++
|
|
return cognitive.AdvisoryResponse{FinishReason: cognitive.FinishCompleted}, nil
|
|
}
|
|
|
|
type admissionTestClient struct {
|
|
decision capability.AdmissionDecision
|
|
releases int
|
|
}
|
|
|
|
func (c *admissionTestClient) Admit(context.Context, string, capability.AdmissionRequest) (capability.AdmissionDecision, capability.AdmissionLease, error) {
|
|
return c.decision, capability.AdmissionLease{ID: "lease-1", PolicyID: "policy-1"}, nil
|
|
}
|
|
func (c *admissionTestClient) Release(context.Context, capability.AdmissionLease) error {
|
|
c.releases++
|
|
return nil
|
|
}
|
|
|
|
func TestInvokeAdmittedDoesNotCallProviderWhenRejected(t *testing.T) {
|
|
provider := &admissionTestProvider{}
|
|
client := &admissionTestClient{decision: capability.AdmissionDecision{State: capability.AdmissionRejected, Reason: capability.AdmissionReasonConcurrency, PolicyID: "policy-1"}}
|
|
gateway := cognitiveGateway{admission: client}
|
|
response, err := gateway.invokeAdmitted(context.Background(), "request-1", provider, cognitive.AdvisoryRequest{Message: "hello"}, 100)
|
|
if err != nil || response.FinishReason != cognitive.FinishRateLimited || provider.calls != 0 || client.releases != 0 {
|
|
t.Fatalf("rejected admission leaked provider call: response=%#v err=%v calls=%d releases=%d", response, err, provider.calls, client.releases)
|
|
}
|
|
}
|
|
|
|
func TestInvokeAdmittedReleasesDurableLeaseAfterProvider(t *testing.T) {
|
|
provider := &admissionTestProvider{}
|
|
client := &admissionTestClient{decision: capability.AdmissionDecision{State: capability.AdmissionAdmitted, Reason: capability.AdmissionReasonNone, PolicyID: "policy-1"}}
|
|
gateway := cognitiveGateway{admission: client}
|
|
response, err := gateway.invokeAdmitted(context.Background(), "request-1", provider, cognitive.AdvisoryRequest{Message: "hello"}, 100)
|
|
if err != nil || response.FinishReason != cognitive.FinishCompleted || provider.calls != 1 || client.releases != 1 {
|
|
t.Fatalf("admitted lease was not released: response=%#v err=%v calls=%d releases=%d", response, err, provider.calls, client.releases)
|
|
}
|
|
}
|