iop/apps/node/internal/transport/session_test.go

package transport_test

import (
	"context"
	"errors"
	"fmt"
	"net"
	"sync"
	"testing"
	"time"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"go.uber.org/zap"
	"google.golang.org/protobuf/proto"

	"iop/apps/node/internal/transport"
	iop "iop/proto/gen/iop"
)

type noopHandler struct{}

func (h *noopHandler) OnRunRequest(_ context.Context, _ *transport.Session, _ *iop.RunRequest) error {
	return nil
}
func (h *noopHandler) OnCancel(_ context.Context, _ *transport.Session, _ *iop.CancelRequest) error {
	return nil
}
func (h *noopHandler) OnCommandRequest(_ context.Context, _ *transport.Session, _ *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error) {
	return nil, nil
}
func (h *noopHandler) OnConfigRefresh(_ context.Context, _ *transport.Session, req *iop.NodeConfigRefreshRequest) (*iop.NodeConfigRefreshResponse, error) {
	return &iop.NodeConfigRefreshResponse{
		RequestId: req.GetRequestId(),
		Status:    iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED,
	}, nil
}

func TestSession_SetHandler_ConcurrentSafe(t *testing.T) {
	var s transport.Session
	var wg sync.WaitGroup
	h := &noopHandler{}
	for i := 0; i < 50; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			s.SetHandler(h)
		}()
	}
	wg.Wait()
}

// buildSessionTestPipe creates a net.Pipe-based pair: one side acts as "edge"
// (sends requests) and the other side acts as the node session under test.
// The edge side parser map must include the response type; the node side must
// include the request type (handled by nodeParserMap via DialEdge, but for
// unit tests we wire it manually).
func buildSessionTestPipe(t *testing.T) (edgeSide *toki.TcpClient, nodeSide *toki.TcpClient) {
	t.Helper()
	edgeConn, nodeConn := net.Pipe()
	edgeParserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.NodeConfigRefreshResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeConfigRefreshResponse{}
			return m, proto.Unmarshal(b, m)
		},
	}
	nodeParserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.NodeConfigRefreshRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeConfigRefreshRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}
	edgeSide = toki.NewTcpClient(edgeConn, 0, 0, edgeParserMap)
	nodeSide = toki.NewTcpClient(nodeConn, 0, 0, nodeParserMap)
	t.Cleanup(func() { edgeSide.Close(); nodeSide.Close() })
	return edgeSide, nodeSide
}

// appliedHandler always returns applied.
type appliedHandler struct{ noopHandler }

func (h *appliedHandler) OnConfigRefresh(_ context.Context, _ *transport.Session, req *iop.NodeConfigRefreshRequest) (*iop.NodeConfigRefreshResponse, error) {
	return &iop.NodeConfigRefreshResponse{
		RequestId: req.GetRequestId(),
		Status:    iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED,
	}, nil
}

// errorHandler returns an error from OnConfigRefresh.
type errorHandler struct{ noopHandler }

func (h *errorHandler) OnConfigRefresh(_ context.Context, _ *transport.Session, req *iop.NodeConfigRefreshRequest) (*iop.NodeConfigRefreshResponse, error) {
	return nil, errors.New("refresh failed")
}

// TestSessionConfigRefreshRequestReturnsHandlerResponse verifies that a
// NodeConfigRefreshRequest pushed by the edge reaches the handler and the
// response is returned to the edge.
func TestSessionConfigRefreshRequestReturnsHandlerResponse(t *testing.T) {
	edgeSide, nodeSide := buildSessionTestPipe(t)

	sess := transport.ExportNewSession(nodeSide, zap.NewNop(), "node-test", "alias-test")
	sess.SetHandler(&appliedHandler{})

	resp, err := toki.SendRequestTyped[*iop.NodeConfigRefreshRequest, *iop.NodeConfigRefreshResponse](
		&edgeSide.Communicator,
		&iop.NodeConfigRefreshRequest{RequestId: "req-1", ChangedPaths: []string{"nodes.0.providers.0.capacity"}},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped: %v", err)
	}
	if resp.GetStatus() != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED {
		t.Fatalf("expected status=applied, got %v", resp.GetStatus())
	}
	if resp.GetRequestId() != "req-1" {
		t.Fatalf("expected request_id=req-1, got %q", resp.GetRequestId())
	}
}

// TestSessionConfigRefreshRequestHandlerErrorReturnsFailure verifies that a
// handler error is translated to a failed protocol response.
func TestSessionConfigRefreshRequestHandlerErrorReturnsFailure(t *testing.T) {
	edgeSide, nodeSide := buildSessionTestPipe(t)

	sess := transport.ExportNewSession(nodeSide, zap.NewNop(), "node-test", "alias-test")
	sess.SetHandler(&errorHandler{})

	resp, err := toki.SendRequestTyped[*iop.NodeConfigRefreshRequest, *iop.NodeConfigRefreshResponse](
		&edgeSide.Communicator,
		&iop.NodeConfigRefreshRequest{RequestId: "req-err"},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped: %v", err)
	}
	if resp.GetStatus() != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_FAILED {
		t.Fatalf("expected status=failed, got %v", resp.GetStatus())
	}
	if resp.GetError() == "" {
		t.Fatal("expected non-empty error message")
	}
}

// TestSessionConfigRefreshRequestNoHandlerReturnsFailure verifies that when no
// handler is set, the session returns a failed response with an informative message.
func TestSessionConfigRefreshRequestNoHandlerReturnsFailure(t *testing.T) {
	edgeSide, nodeSide := buildSessionTestPipe(t)

	_ = transport.ExportNewSession(nodeSide, zap.NewNop(), "node-nohandler", "")
	// handler intentionally NOT set

	resp, err := toki.SendRequestTyped[*iop.NodeConfigRefreshRequest, *iop.NodeConfigRefreshResponse](
		&edgeSide.Communicator,
		&iop.NodeConfigRefreshRequest{RequestId: "req-nohandler"},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped: %v", err)
	}
	if resp.GetStatus() != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_FAILED {
		t.Fatalf("expected status=failed when handler is nil, got %v", resp.GetStatus())
	}
}

// Compile check: Session must export a way to create instances for tests.
// ExportNewSession is expected in session_export_test.go or a separate test helper file.
var _ = fmt.Sprintf