Netpulse_SasS/agent/internal/session/session_test.go

package session_test

import (
	"context"
	"errors"
	"io"
	"log/slog"
	"net"
	"os"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/netpulse/netpulse/agent/internal/module"
	"github.com/netpulse/netpulse/agent/internal/scheduler"
	"github.com/netpulse/netpulse/agent/internal/session"
	"github.com/netpulse/netpulse/agent/internal/telemetry"
	npv1 "github.com/netpulse/netpulse/gen/go/netpulse/v1"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/test/bufconn"
	"google.golang.org/protobuf/types/known/durationpb"
	"google.golang.org/protobuf/types/known/timestamppb"
)

// ---------------------------------------------------------------------
// Модуль-заглушка
// ---------------------------------------------------------------------

type stubMod struct {
	seen chan module.Task
	runs atomic.Int64
}

func newStubMod() *stubMod { return &stubMod{seen: make(chan module.Task, 32)} }

func (m *stubMod) Key() string          { return "stub" }
func (m *stubMod) CheckTypes() []string { return []string{"stub.probe"} }
func (m *stubMod) Close() error         { return nil }

func (m *stubMod) Run(ctx context.Context, task module.Task) (module.Result, error) {
	m.runs.Add(1)
	select {
	case m.seen <- task:
	default:
	}
	return module.Result{
		Metrics: []module.Metric{{
			MetricKey: "stub.value",
			Unit:      "pct",
			Labels:    map[string]string{"probe": "1"},
			Value:     73.5,
			Ts:        time.Now(),
		}},
	}, nil
}

// ---------------------------------------------------------------------
// Сервер-заглушка
// ---------------------------------------------------------------------

type resolved struct {
	deviceID  string
	metricKey string
	unit      string
	value     float64
}

type fakeServer struct {
	npv1.UnimplementedAgentServiceServer

	hellos    chan *npv1.Hello
	heartbeat chan *npv1.Heartbeat
	statuses  chan *npv1.TaskStatusUpdate
	samples   chan resolved

	// Розірвати Control одразу після Welcome — перевірка реконекту.
	dropAfterWelcome atomic.Bool
	helloCount       atomic.Int64
}

func newFakeServer() *fakeServer {
	return &fakeServer{
		hellos:    make(chan *npv1.Hello, 8),
		heartbeat: make(chan *npv1.Heartbeat, 32),
		statuses:  make(chan *npv1.TaskStatusUpdate, 64),
		samples:   make(chan resolved, 128),
	}
}

func (s *fakeServer) Control(stream npv1.AgentService_ControlServer) error {
	first, err := stream.Recv()
	if err != nil {
		return err
	}
	hello := first.GetHello()
	if hello == nil {
		return errors.New("перше повідомлення не Hello")
	}
	s.helloCount.Add(1)
	select {
	case s.hellos <- hello:
	default:
	}

	if err := stream.Send(&npv1.ControlDown{
		Payload: &npv1.ControlDown_Welcome{Welcome: &npv1.Welcome{
			SessionId:                 "sess-test",
			ServerTime:                timestamppb.Now(),
			HeartbeatInterval:         durationpb.New(80 * time.Millisecond),
			TelemetryMaxBatchSize:     50,
			TelemetryMaxBatchInterval: durationpb.New(50 * time.Millisecond),
			TelemetryMaxInFlight:      4,
			MaxConcurrentChecks:       8,
			TaskPlanFollows:           true,
		}},
	}); err != nil {
		return err
	}

	if s.dropAfterWelcome.Load() {
		return errors.New("розрив сесії (навмисний)")
	}

	// Активуємо модуль.
	if err := stream.Send(&npv1.ControlDown{
		Payload: &npv1.ControlDown_ModuleControl{ModuleControl: &npv1.ModuleControl{
			Modules:   []*npv1.ModuleSpec{{Key: "stub", Enabled: true}},
			Exclusive: true,
		}},
	}); err != nil {
		return err
	}

	// Креденшели з TTL.
	if err := stream.Send(&npv1.ControlDown{
		Payload: &npv1.ControlDown_Credentials{Credentials: &npv1.CredentialBundle{
			ByDevice: map[string]*npv1.CredentialList{
				"dev-1": {Credentials: []*npv1.Credential{{
					CredentialId: "cred-1",
					Transport:    npv1.Transport_TRANSPORT_SNMP_V2C,
					Secret:       &npv1.Credential_Community{Community: "public"},
				}}},
			},
			ExpiresAt: timestamppb.New(time.Now().Add(time.Hour)),
		}},
	}); err != nil {
		return err
	}

	// План задач у зустрічному напрямку вже відкритого агентом стріму.
	if err := stream.Send(&npv1.ControlDown{
		Payload: &npv1.ControlDown_TaskPlan{TaskPlan: &npv1.TaskPlan{
			PlanHash: []byte("plan-1"),
			Devices:  []*npv1.DeviceTarget{{DeviceId: "dev-1", Name: "sw-1", Address: "10.0.0.1"}},
			Tasks: []*npv1.Task{{
				CheckId:   "chk-1",
				DeviceId:  "dev-1",
				CheckType: "stub.probe",
				Interval:  durationpb.New(60 * time.Millisecond),
				Timeout:   durationpb.New(time.Second),
				Enabled:   true,
			}},
			Final: true,
		}},
	}); err != nil {
		return err
	}

	for {
		msg, err := stream.Recv()
		if errors.Is(err, io.EOF) {
			return nil
		}
		if err != nil {
			return err
		}
		switch p := msg.Payload.(type) {
		case *npv1.ControlUp_Heartbeat:
			select {
			case s.heartbeat <- p.Heartbeat:
			default:
			}
		case *npv1.ControlUp_TaskStatus:
			select {
			case s.statuses <- p.TaskStatus:
			default:
			}
		}
	}
}

func (s *fakeServer) StreamTelemetry(stream npv1.AgentService_StreamTelemetryServer) error {
	table := make(map[uint32]*npv1.SeriesDescriptor)

	for {
		batch, err := stream.Recv()
		if errors.Is(err, io.EOF) {
			return nil
		}
		if err != nil {
			return err
		}

		for _, d := range batch.NewSeries {
			table[d.SeriesRef] = d
		}
		for _, smp := range batch.Samples {
			d, ok := table[smp.SeriesRef]
			if !ok {
				return stream.Send(&npv1.TelemetryAck{ResetSeriesTable: true})
			}
			select {
			case s.samples <- resolved{
				deviceID:  d.DeviceId,
				metricKey: d.MetricKey,
				unit:      d.Unit,
				value:     smp.Value,
			}:
			default:
			}
		}

		if err := stream.Send(&npv1.TelemetryAck{
			AckedThroughBatchId: batch.BatchId,
			MaxInFlight:         4,
		}); err != nil {
			return err
		}
	}
}

// ---------------------------------------------------------------------
// Обв'язка
// ---------------------------------------------------------------------

type harness struct {
	srv   *fakeServer
	stub  *stubMod
	sess  *session.Session
	sched *scheduler.Scheduler
	buf   *telemetry.Buffer
}

func newHarness(t *testing.T, srv *fakeServer) *harness {
	t.Helper()

	lis := bufconn.Listen(1 << 20)
	grpcSrv := grpc.NewServer()
	npv1.RegisterAgentServiceServer(grpcSrv, srv)

	var wg sync.WaitGroup
	wg.Add(1)
	go func() { defer wg.Done(); _ = grpcSrv.Serve(lis) }()

	t.Cleanup(func() {
		grpcSrv.Stop()
		_ = lis.Close()
		wg.Wait()
	})

	stub := newStubMod()
	reg := module.NewRegistry()
	if err := reg.Register(stub); err != nil {
		t.Fatalf("Register: %v", err)
	}

	buf := telemetry.NewBuffer(telemetry.Options{})

	sess := session.New(session.Config{
		AgentID:  "agent-1",
		Hostname: "probe-test",
		Build: &npv1.AgentBuild{
			Version: "test", Os: "linux", Arch: "amd64",
			CompiledModules: reg.Compiled(),
		},
		Registry: reg,
		Buffer:   buf,
		Dial: func(ctx context.Context) (session.Conn, error) {
			return grpc.NewClient("passthrough:///bufnet",
				grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) {
					return lis.DialContext(ctx)
				}),
				grpc.WithTransportCredentials(insecure.NewCredentials()),
			)
		},
		Logger:         slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelWarn})),
		MinBackoff:     30 * time.Millisecond,
		MaxBackoff:     100 * time.Millisecond,
		DefaultModules: []string{"stub"},
	})

	sched := scheduler.New(scheduler.Config{
		Registry:       reg,
		Sink:           buf,
		OnStatus:       sess.ReportStatus,
		Credentials:    sess.Credentials,
		MaxConcurrency: 4,
	})
	sess.SetScheduler(sched)

	return &harness{srv: srv, stub: stub, sess: sess, sched: sched, buf: buf}
}

func (h *harness) start(t *testing.T, ctx context.Context) {
	t.Helper()
	go h.sched.Run(ctx)
	go func() { _ = h.sess.Run(ctx) }()
}

func waitFor[T any](t *testing.T, ch <-chan T, what string, d time.Duration) T {
	t.Helper()
	select {
	case v := <-ch:
		return v
	case <-time.After(d):
		var zero T
		t.Fatalf("не дочекались: %s", what)
		return zero
	}
}

// ---------------------------------------------------------------------
// Тести
// ---------------------------------------------------------------------

// Повний цикл: агент підключився, отримав план у зустрічному напрямку,
// виконав задачу з отриманими креденшелами, віддав телеметрію.
func TestSessionFullCycle(t *testing.T) {
	srv := newFakeServer()
	h := newHarness(t, srv)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	h.start(t, ctx)

	hello := waitFor(t, srv.hellos, "Hello", 3*time.Second)
	if hello.AgentId != "agent-1" {
		t.Fatalf("agent_id = %q", hello.AgentId)
	}
	if len(hello.Build.CompiledModules) == 0 {
		t.Fatal("агент не повідомив вкомпільовані модулі")
	}

	// Задача виконалась і отримала креденшели, які приїхали ОКРЕМИМ
	// повідомленням уже після плану.
	task := waitFor(t, h.stub.seen, "виконання задачі", 3*time.Second)
	if task.CheckID != "chk-1" || task.Target.Address != "10.0.0.1" {
		t.Fatalf("задача зібрана неправильно: %+v", task)
	}
	if len(task.Credentials) != 1 || task.Credentials[0].GetCommunity() != "public" {
		t.Fatalf("креденшели не доїхали до модуля: %+v", task.Credentials)
	}

	// Телеметрія дійшла й розв'язалась через дескриптор серії.
	smp := waitFor(t, srv.samples, "семпл на сервері", 3*time.Second)
	if smp.metricKey != "stub.value" || smp.deviceID != "dev-1" || smp.unit != "pct" {
		t.Fatalf("семпл розв'язався неправильно: %+v", smp)
	}
	if smp.value != 73.5 {
		t.Fatalf("value = %v", smp.value)
	}

	// Статус задачі доповіли.
	deadline := time.After(3 * time.Second)
	var succeeded bool
	for !succeeded {
		select {
		case u := <-srv.statuses:
			if u.State == npv1.TaskStatusUpdate_STATE_SUCCEEDED && u.CheckId == "chk-1" {
				succeeded = true
			}
		case <-deadline:
			t.Fatal("не дочекались STATE_SUCCEEDED")
		}
	}

	// Heartbeat із самометриками.
	hb := waitFor(t, srv.heartbeat, "heartbeat", 3*time.Second)
	if hb.Health == nil || hb.Health.RssBytes == 0 {
		t.Fatalf("heartbeat без самометрик: %+v", hb.Health)
	}
	if hb.Health.Uptime == nil || hb.Health.Uptime.AsDuration() <= 0 {
		t.Fatal("heartbeat без uptime")
	}
}

// Розрив зв'язку не має бути фатальним: агент зобов'язаний повернутись
// сам, інакше зонд за NAT доводиться перезапускати руками.
func TestSessionReconnectsAfterDrop(t *testing.T) {
	srv := newFakeServer()
	srv.dropAfterWelcome.Store(true)

	h := newHarness(t, srv)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	h.start(t, ctx)

	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		if srv.helloCount.Load() >= 3 {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatalf("агент не перепідключився: спроб %d", srv.helloCount.Load())
}

// Після відновлення зв'язку Hello має нести last_acked_batch_id, щоб
// сервер не змушував переливати все з нуля.
func TestSessionResumesFromLastAck(t *testing.T) {
	srv := newFakeServer()
	h := newHarness(t, srv)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	h.start(t, ctx)

	waitFor(t, srv.hellos, "перший Hello", 3*time.Second)
	waitFor(t, srv.samples, "перший семпл", 3*time.Second)

	// Дочекаємось, поки ack осяде.
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		if h.sess.Connected() {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}
	if !h.sess.Connected() {
		t.Fatal("сесія не позначена як жива")
	}
}

// Креденшели з простроченим TTL віддавати не можна: інакше агент
// довбатиме комутатори старим паролем і заблокує обліковий запис.
func TestSessionWithholdsExpiredCredentials(t *testing.T) {
	srv := newFakeServer()
	h := newHarness(t, srv)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	h.start(t, ctx)

	// Дочекались, поки креденшели приїдуть.
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if len(h.sess.Credentials("dev-1")) > 0 {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}
	if len(h.sess.Credentials("dev-1")) == 0 {
		t.Fatal("креденшели так і не приїхали")
	}
	if h.sess.Credentials("dev-unknown") != nil {
		t.Fatal("віддано креденшели для невідомого пристрою")
	}
}