File Coverage

lib/PAGI/Test/WebSocket.pm

Criterion	Covered	Total	%
statement	116	124	93.5
branch	22	36	61.1
condition	16	28	57.1
subroutine	20	20	100.0
pod	11	11	100.0
total	185	219	84.4

package PAGI::Test::WebSocket;

use strict;
use warnings;
use Future::AsyncAwait;
use Future;
use Carp qw(croak);


sub new {
    my ($class, %args) = @_;

    croak "app is required" unless $args{app};
    croak "scope is required" unless $args{scope};

    return bless {
        app         => $args{app},
        scope       => $args{scope},
        send_queue  => [],      # Messages from test -> app
        recv_queue  => [],      # Messages from app -> test
        closed      => 0,
        accepted    => 0,
        close_code  => undef,
        close_reason => '',
        _pending_receives => [],  # Pending receive futures
    }, $class;
}

sub _start {
    my ($self) = @_;

    # Create receive coderef for the app
    my $receive = async sub {
        # First call returns websocket.connect
        if (!$self->{_connect_sent}) {
            $self->{_connect_sent} = 1;
            return { type => 'websocket.connect' };
        }

        # Return queued message if available
        if (@{$self->{send_queue}}) {
            return shift @{$self->{send_queue}};
        }

        # Return disconnect if closed
        if ($self->{closed}) {
            return { type => 'websocket.disconnect', code => $self->{close_code} // 1000 };
        }

        # Create a future that will be resolved when data arrives
        my $future = Future->new;
        push @{$self->{_pending_receives}}, $future;
        return await $future;
    };

    # Create send coderef for the app
    my $send = async sub {
        my ($event) = @_;

        if ($event->{type} eq 'websocket.accept') {
            $self->{accepted} = 1;
        }
        elsif ($event->{type} eq 'websocket.send') {
            push @{$self->{recv_queue}}, $event;
        }
        elsif ($event->{type} eq 'websocket.close') {
            $self->{closed} = 1;
            $self->{close_code} = $event->{code} // 1000;
            $self->{close_reason} = $event->{reason} // '';
        }

        return;
    };

    # Start the app future but don't block on it
    $self->{app_future} = $self->{app}->($self->{scope}, $receive, $send);

    # Wait for acceptance (the first two awaits in the app should complete immediately)
    # This is a bit hacky but works: we need to let the app run until it accepts
    $self->_pump_app;

    croak "WebSocket connection not accepted" unless $self->{accepted};

    return $self;
}

sub _pump_app {
    my ($self) = @_;

    # This pumps the app future by checking if it's waiting on a receive
    # If there are pending receives and we have data, resolve them
    while (@{$self->{_pending_receives}} && @{$self->{send_queue}}) {
        my $future = shift @{$self->{_pending_receives}};
        my $event = shift @{$self->{send_queue}};
        $future->done($event);
    }

    # If closed and there are pending receives, resolve them with disconnect
    if ($self->{closed} && @{$self->{_pending_receives}}) {
        while (my $future = shift @{$self->{_pending_receives}}) {
            $future->done({ type => 'websocket.disconnect', code => $self->{close_code} // 1000 });
        }
    }
}

sub send_text {
    my ($self, $text) = @_;

    croak "Cannot send on closed WebSocket" if $self->{closed};

    push @{$self->{send_queue}}, {
        type => 'websocket.receive',
        text => $text,
    };

    # Pump the app to process this message
    $self->_pump_app;

    return $self;
}

sub send_bytes {
    my ($self, $bytes) = @_;

    croak "Cannot send on closed WebSocket" if $self->{closed};

    push @{$self->{send_queue}}, {
        type => 'websocket.receive',
        bytes => $bytes,
    };

    # Pump the app to process this message
    $self->_pump_app;

    return $self;
}

sub send_json {
    my ($self, $data) = @_;

    require JSON::MaybeXS;
    my $text = JSON::MaybeXS::encode_json($data);

    return $self->send_text($text);
}

sub receive_text {
    my ($self, $timeout) = @_;
    $timeout //= 5;

    # Check if we have a text message already waiting
    for my $i (0 .. $#{$self->{recv_queue}}) {
        my $event = $self->{recv_queue}[$i];
        if ($event->{type} eq 'websocket.send' && exists $event->{text}) {
            splice @{$self->{recv_queue}}, $i, 1;
            return $event->{text};
        }
    }

    # Check if connection closed
    return undef if $self->{closed};

    # No message available yet
    croak "Timeout waiting for WebSocket text message";
}

sub receive_bytes {
    my ($self, $timeout) = @_;
    $timeout //= 5;

    # Check if we have a bytes message waiting
    for my $i (0 .. $#{$self->{recv_queue}}) {
        my $event = $self->{recv_queue}[$i];
        if ($event->{type} eq 'websocket.send' && exists $event->{bytes}) {
            splice @{$self->{recv_queue}}, $i, 1;
            return $event->{bytes};
        }
    }

    # Check if connection closed
    return undef if $self->{closed};

    # No message available yet
    croak "Timeout waiting for WebSocket bytes message";
}

sub receive_json {
    my ($self, $timeout) = @_;

    my $text = $self->receive_text($timeout);
    return undef unless defined $text;

    require JSON::MaybeXS;
    return JSON::MaybeXS::decode_json($text);
}

sub close {
    my ($self, $code, $reason) = @_;

    return if $self->{closed};

    $code //= 1000;
    $reason //= '';

    $self->{closed} = 1;
    $self->{close_code} = $code;
    $self->{close_reason} = $reason;

    # Push disconnect event
    push @{$self->{send_queue}}, {
        type => 'websocket.disconnect',
        code => $code,
        reason => $reason,
    };

    # Pump the app to let it process the disconnect
    $self->_pump_app;

    return $self;
}

sub close_code {
    my ($self) = @_;
    return $self->{close_code};
}

sub close_reason {
    my ($self) = @_;
    return $self->{close_reason};
}

sub is_closed {
    my ($self) = @_;
    return $self->{closed};
}

1;

__END__