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