mirror of
https://github.com/Buriburizaem0n/nezha_domains.git
synced 2026-09-22 03:00:12 +00:00
test(agentcompat): add integration scenarios
Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>
This commit is contained in:
naiba
naiba/CloudCode
@@ -0,0 +1,233 @@
|
||||
//go:build linux
|
||||
|
||||
package scenario
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"reflect"
|
||||
|
||||
"github.com/nezhahq/nezha/integration/agentcompat/internal/agent"
|
||||
"github.com/nezhahq/nezha/integration/agentcompat/internal/client"
|
||||
"github.com/nezhahq/nezha/integration/agentcompat/internal/contract"
|
||||
"github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrInvalidHeldSessionSetTopology = errors.New("held session set topology is invalid")
|
||||
ErrInvalidHeldSessionSetPlan = errors.New("held session set plan is invalid")
|
||||
)
|
||||
|
||||
const heldSessionSetAgentCount = 8
|
||||
|
||||
type HeldSessionAgent struct {
|
||||
Ordinal StressAgentOrdinal
|
||||
Agent *agent.Agent
|
||||
Readiness agent.Readiness
|
||||
PATClient *client.Client
|
||||
}
|
||||
|
||||
type HeldSessionSetInput struct {
|
||||
Dashboard *dashboard.Dashboard
|
||||
Plan StressPlan
|
||||
Topology []HeldSessionAgent
|
||||
Dependencies HeldSessionSetDependencies
|
||||
ControlClient *client.Client
|
||||
ControlServerIDs []uint64
|
||||
testHealthSnapshotRequestHook func(int)
|
||||
testHealthSnapshotReplyHook func(int)
|
||||
testHealthSnapshotOverrideHook func(int, heldHealthSnapshotRequest) *heldHealthSnapshot
|
||||
testHealthSnapshotSendHook func(int)
|
||||
testHealthEventHook func(heldHealthMessage)
|
||||
testHealthClosureObservedHook func(int)
|
||||
testHealthSnapshotAcceptedHook func(heldHealthSnapshot)
|
||||
testHealthShutdownAcceptedHook func()
|
||||
testHealthShutdownAcknowledgedHook func()
|
||||
}
|
||||
|
||||
type heldSessionSetTopology struct {
|
||||
dashboard *dashboard.Dashboard
|
||||
stateClient heldSessionSetStateObserver
|
||||
agents map[int]HeldSessionAgent
|
||||
}
|
||||
|
||||
type heldSessionSetStateObserver interface {
|
||||
IOStreamState(context.Context) (client.IOStreamState, error)
|
||||
WaitForIOStreamState(context.Context, client.IOStreamStateExpectation) (client.IOStreamState, error)
|
||||
}
|
||||
|
||||
type heldSessionSetAuthorizedObserver struct {
|
||||
client *client.Client
|
||||
}
|
||||
|
||||
func (observer heldSessionSetAuthorizedObserver) IOStreamState(ctx context.Context) (client.IOStreamState, error) {
|
||||
return observer.client.IOStreamState(ctx)
|
||||
}
|
||||
|
||||
func (observer heldSessionSetAuthorizedObserver) WaitForIOStreamState(ctx context.Context, expectation client.IOStreamStateExpectation) (client.IOStreamState, error) {
|
||||
return observer.client.WaitForIOStreamState(ctx, expectation)
|
||||
}
|
||||
|
||||
type HeldSessionSetDependencies struct {
|
||||
Terminal func(context.Context, heldTerminalInput) (heldSession, error)
|
||||
NAT func(context.Context, heldNATInput) (heldSession, error)
|
||||
FM func(context.Context, heldLegacyFMInput) (heldSession, error)
|
||||
Snapshot func(context.Context, heldSessionSetStateObserver) (client.IOStreamState, error)
|
||||
WaitState func(context.Context, heldSessionSetStateObserver, client.IOStreamStateExpectation) (client.IOStreamState, error)
|
||||
InspectAgent func(*agent.Agent) heldSessionAgentFacts
|
||||
ObserveState func(*client.Client) heldSessionSetStateObserver
|
||||
}
|
||||
|
||||
type heldSessionAgentFacts struct {
|
||||
PID int
|
||||
UUID string
|
||||
}
|
||||
|
||||
func defaultHeldSessionSetDependencies() HeldSessionSetDependencies {
|
||||
return HeldSessionSetDependencies{
|
||||
Terminal: func(ctx context.Context, input heldTerminalInput) (heldSession, error) {
|
||||
return newHeldTerminalSession(ctx, input)
|
||||
},
|
||||
NAT: func(ctx context.Context, input heldNATInput) (heldSession, error) {
|
||||
return newHeldNATSession(ctx, input)
|
||||
},
|
||||
FM: func(ctx context.Context, input heldLegacyFMInput) (heldSession, error) {
|
||||
return newHeldLegacyFMSession(ctx, input)
|
||||
},
|
||||
Snapshot: func(ctx context.Context, stateClient heldSessionSetStateObserver) (client.IOStreamState, error) {
|
||||
return stateClient.IOStreamState(ctx)
|
||||
},
|
||||
WaitState: func(ctx context.Context, stateClient heldSessionSetStateObserver, expectation client.IOStreamStateExpectation) (client.IOStreamState, error) {
|
||||
return stateClient.WaitForIOStreamState(ctx, expectation)
|
||||
},
|
||||
InspectAgent: func(instance *agent.Agent) heldSessionAgentFacts {
|
||||
return heldSessionAgentFacts{PID: instance.PID(), UUID: instance.UUID()}
|
||||
},
|
||||
ObserveState: func(controlClient *client.Client) heldSessionSetStateObserver {
|
||||
return heldSessionSetAuthorizedObserver{client: controlClient}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func validateHeldSessionSetPlans(plan StressPlan) ([]StressSessionPlan, error) {
|
||||
canonical, err := canonicalHeldSessionPlans(plan)
|
||||
if err != nil {
|
||||
return nil, errors.Join(ErrInvalidHeldSessionSetPlan, err)
|
||||
}
|
||||
if !reflect.DeepEqual(canonical, plan.Sessions) || len(plan.Sessions) != 12 {
|
||||
return nil, ErrInvalidHeldSessionSetPlan
|
||||
}
|
||||
ids := make(map[string]struct{}, len(plan.Sessions))
|
||||
for _, session := range plan.Sessions {
|
||||
if session.ID.String() == "" {
|
||||
return nil, fmt.Errorf("empty session plan ID: %w", ErrInvalidHeldSessionSetPlan)
|
||||
}
|
||||
if _, exists := ids[session.ID.String()]; exists {
|
||||
return nil, fmt.Errorf("duplicate session plan ID: %s: %w", session.ID.String(), ErrInvalidHeldSessionSetPlan)
|
||||
}
|
||||
ids[session.ID.String()] = struct{}{}
|
||||
}
|
||||
if countHeldSessionKind(plan.Sessions, StressSessionTerminal) != 4 || countHeldSessionKind(plan.Sessions, StressSessionNAT) != 4 || countHeldSessionKind(plan.Sessions, StressSessionFM) != 4 {
|
||||
return nil, ErrInvalidHeldSessionSetPlan
|
||||
}
|
||||
return append([]StressSessionPlan(nil), plan.Sessions...), nil
|
||||
}
|
||||
|
||||
func canonicalHeldSessionPlans(plan StressPlan) ([]StressSessionPlan, error) {
|
||||
if plan.Seed == 0 || plan.Profile == "" {
|
||||
return nil, ErrInvalidHeldSessionSetPlan
|
||||
}
|
||||
profile, err := contract.ProfileByName(string(plan.Profile))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
canonical, err := GenerateStressPlan(profile, plan.Seed)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return canonical.Sessions, nil
|
||||
}
|
||||
|
||||
func countHeldSessionKind(plans []StressSessionPlan, kind StressSessionKind) int {
|
||||
count := 0
|
||||
for _, plan := range plans {
|
||||
if plan.Kind == kind {
|
||||
count++
|
||||
}
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func validateHeldSessionSetTopology(input HeldSessionSetInput, plans []StressSessionPlan) (heldSessionSetTopology, error) {
|
||||
if input.Dashboard == nil || input.ControlClient == nil {
|
||||
return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology
|
||||
}
|
||||
inspectAgent := input.Dependencies.InspectAgent
|
||||
if inspectAgent == nil {
|
||||
inspectAgent = defaultHeldSessionSetDependencies().InspectAgent
|
||||
}
|
||||
observeState := input.Dependencies.ObserveState
|
||||
if observeState == nil {
|
||||
observeState = defaultHeldSessionSetDependencies().ObserveState
|
||||
}
|
||||
profile, err := contract.ProfileByName(string(input.Plan.Profile))
|
||||
if err != nil || profile.AgentCount() != heldSessionSetAgentCount || len(input.Topology) != heldSessionSetAgentCount {
|
||||
return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology
|
||||
}
|
||||
agents := make(map[int]HeldSessionAgent, len(input.Topology))
|
||||
uuids := make(map[string]struct{}, len(input.Topology))
|
||||
serverIDs := make(map[uint64]struct{}, len(input.Topology))
|
||||
for _, topology := range input.Topology {
|
||||
ordinal := topology.Ordinal.Int()
|
||||
if ordinal < 1 || ordinal > heldSessionSetAgentCount || topology.Agent == nil || topology.PATClient == nil {
|
||||
return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology
|
||||
}
|
||||
if _, exists := agents[ordinal]; exists {
|
||||
return heldSessionSetTopology{}, fmt.Errorf("duplicate agent ordinal %d: %w", ordinal, ErrInvalidHeldSessionSetTopology)
|
||||
}
|
||||
agentFacts := inspectAgent(topology.Agent)
|
||||
if err := validateHeldReadinessFacts(agentFacts, topology.Readiness); err != nil {
|
||||
return heldSessionSetTopology{}, err
|
||||
}
|
||||
if agentFacts.PID < 1 || topology.Readiness.ServerID == 0 || topology.Readiness.UUID != agentFacts.UUID {
|
||||
return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology
|
||||
}
|
||||
if _, exists := uuids[topology.Readiness.UUID]; exists {
|
||||
return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology
|
||||
}
|
||||
if _, exists := serverIDs[topology.Readiness.ServerID]; exists {
|
||||
return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology
|
||||
}
|
||||
uuids[topology.Readiness.UUID] = struct{}{}
|
||||
serverIDs[topology.Readiness.ServerID] = struct{}{}
|
||||
agents[ordinal] = topology
|
||||
}
|
||||
if len(input.ControlServerIDs) != heldSessionSetAgentCount {
|
||||
return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology
|
||||
}
|
||||
controlServerIDs := make(map[uint64]struct{}, len(input.ControlServerIDs))
|
||||
for _, serverID := range input.ControlServerIDs {
|
||||
if serverID == 0 {
|
||||
return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology
|
||||
}
|
||||
if _, exists := serverIDs[serverID]; !exists {
|
||||
return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology
|
||||
}
|
||||
if _, exists := controlServerIDs[serverID]; exists {
|
||||
return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology
|
||||
}
|
||||
controlServerIDs[serverID] = struct{}{}
|
||||
}
|
||||
for ordinal := 1; ordinal <= heldSessionSetAgentCount; ordinal++ {
|
||||
if _, exists := agents[ordinal]; !exists {
|
||||
return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology
|
||||
}
|
||||
}
|
||||
for _, plan := range plans {
|
||||
if _, exists := agents[plan.Agent.Int()]; !exists {
|
||||
return heldSessionSetTopology{}, fmt.Errorf("missing agent ordinal %d: %w", plan.Agent.Int(), ErrInvalidHeldSessionSetTopology)
|
||||
}
|
||||
}
|
||||
return heldSessionSetTopology{dashboard: input.Dashboard, stateClient: observeState(input.ControlClient), agents: agents}, nil
|
||||
}
|
||||
Reference in New Issue
Block a user