#!/usr/bin/perl
#
#  Copyright (c) 2011-2017 FastMail Pty Ltd. All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions
#  are met:
#
#  1. Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
#
#  2. Redistributions in binary form must reproduce the above copyright
#     notice, this list of conditions and the following disclaimer in
#     the documentation and/or other materials provided with the
#     distribution.
#
#  3. The name "Fastmail Pty Ltd" must not be used to
#     endorse or promote products derived from this software without
#     prior written permission. For permission or any legal
#     details, please contact
#      FastMail Pty Ltd
#      PO Box 234
#      Collins St West 8007
#      Victoria
#      Australia
#
#  4. Redistributions of any form whatsoever must retain the following
#     acknowledgment:
#     "This product includes software developed by Fastmail Pty. Ltd."
#
#  FASTMAIL PTY LTD DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE,
#  INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY  AND FITNESS, IN NO
#  EVENT SHALL OPERA SOFTWARE AUSTRALIA BE LIABLE FOR ANY SPECIAL, INDIRECT
#  OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF
#  USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER
#  TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE
#  OF THIS SOFTWARE.
#

package Cassandane::Unit::TestPlanItem;
use strict;
use warnings;
use IO::Handle;
use POSIX;
use Time::HiRes qw(time);

use lib '.';
use Cassandane::Unit::TestCase;

sub new
{
    my ($class, $suite) = @_;
    my $self = {
	suite => $suite,
	loaded_suite => undef,
	denied => {},
	allowed => {},
    };
    return bless $self, $class;
}

sub _get_loaded_suite
{
    my ($self) = @_;
    return $self->{loaded_suite} ||= Test::Unit::Loader::load($self->{suite});
}

sub _is_allowed
{
    my ($self, $name) = @_;

    # Rules are:
    # deny if method has been explicitly denied
    return 0 if $self->{denied}->{$name};
    # allow if method has been explicitly allowed
    return 1 if $self->{allowed}->{$name};
    # deny if anything is explicitly allowed
    return 0 if scalar keys %{$self->{allowed}};
    # finally, allow
    return 1;
}

sub _deny
{
    my ($self, $name) = @_;
    $self->{denied}->{$name} = 1;
}

sub _allow
{
    my ($self, $name) = @_;
    $self->{allowed}->{$name} = 1;
}

package Cassandane::Unit::Worker;
use Storable qw(freeze thaw);
use MIME::Base64;

my $nextid = 1;

sub new
{
    my ($class) = @_;
    my $self = {
	id => $nextid++,
	pid => undef,
	downpipe => undef,
	uppipe => undef,
	busy => 0,
	handler => undef,
    };
    return bless $self, $class;
}

sub _pipe_read_fh
{
    my ($r, $w) = @_;

    POSIX::close($w);
    my $fh = IO::Handle->new_from_fd($r, "r");
    $fh->autoflush(1);
    return $fh;
}

sub _pipe_write_fh
{
    my ($r, $w) = @_;

    POSIX::close($r);
    my $fh = IO::Handle->new_from_fd($w, "w");
    $fh->autoflush(1);
    return $fh;
}

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

    my ($dr, $dw) = POSIX::pipe();
    die "Cannot create down pipe: $!"
	unless defined $dw;

    my ($ur, $uw) = POSIX::pipe();
    die "Cannot create up pipe: $!"
	unless defined $uw;

    my $pid = fork();
    die "Cannot fork: $!" unless defined $pid;

    if ($pid)
    {
	# parent
	$self->{downpipe} = _pipe_write_fh($dr, $dw);
	$self->{uppipe} = _pipe_read_fh($ur, $uw);
	$self->{pid} = $pid;
    }
    else
    {
	# child
	$self->{downpipe} = _pipe_read_fh($dr, $dw);
	$self->{uppipe} = _pipe_write_fh($ur, $uw);
	$ENV{TEST_UNIT_WORKER_ID} = $self->{id};    # 1, 2, 3...
	$ENV{TEST_UNIT_BASENAME} = $0;
	$0 = "$ENV{TEST_UNIT_BASENAME} ($ENV{TEST_UNIT_WORKER_ID})";
	$self->_mainloop();
	exit(0);
    }
}

sub _send
{
    my ($fh, $fmt, @args) = @_;
    my $msg = sprintf($fmt, @args);
# print STDERR "--> \"$msg\"\n";
    syswrite($fh, $msg)
	or die "Cannot write to pipe: $!";
}

sub _receive
{
    my ($fh) = @_;
    my $msg = $fh->gets()
	or return;
# print STDERR "<-- \"$msg\"\n";
    chomp $msg;
    return $msg;
}

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

    while (my $msg = _receive($self->{downpipe}))
    {
	my ($command, @args) = split(/\s+/, $msg);

	if ($command eq 'stop')
	{
	    return;
	}
	elsif ($command eq 'run')
	{
	    my ($witem) = thaw(decode_base64($args[0]));
	    $0 = "$ENV{TEST_UNIT_BASENAME} ($ENV{TEST_UNIT_WORKER_ID}) $witem->{suite}.$witem->{testname}";
	    $self->{handler}->($witem);
	    $0 = "$ENV{TEST_UNIT_BASENAME} ($ENV{TEST_UNIT_WORKER_ID})";
	    _send($self->{uppipe},
		  "done %s\n", encode_base64(freeze($witem), ''));
	}
	else
	{
	    print STDERR "_mainloop: unknown command '$command'\n";
	}
    }
}

sub get_reply
{
    my ($self) = @_;
    return if !$self->{busy};
    my $msg = _receive($self->{uppipe});
    return if !defined $msg;
    my ($command, @args) = split(/\s+/, $msg);
    die "Unknown message \"$msg\""
	if ($command ne 'done');
    $self->{busy} = 0;
    my ($witem) = thaw(decode_base64($args[0]));
    return $witem;
}

sub assign
{
    my ($self, $witem) = @_;
    $witem->{start_time} = time();
    _send($self->{downpipe},
	  "run %s\n", encode_base64(freeze($witem), ''));
    $self->{busy} = 1;
}

sub stop
{
    my ($self) = @_;
    eval
    {
	# We don't care if this dies, it just
	# means the Worker has died prematurely.
	_send($self->{downpipe}, "stop\n");
    };
    while (1)
    {
	my $res = waitpid($self->{pid}, 0);
	last if ($res < 0 || $res == $self->{pid});
    }
    $self->_cleanup();
}

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

    if ($self->{downpipe})
    {
	close $self->{downpipe};
	$self->{downpipe} = undef;
    }

    if ($self->{uppipe})
    {
	close $self->{uppipe};
	$self->{uppipe} = undef;
    }

    $self->{pid} = undef;
}

sub DESTROY
{
    my ($self) = @_;
    $self->_cleanup();
}

package Cassandane::Unit::WorkerPool;

use Errno qw(EINTR);

sub new
{
    my ($class, %params) = @_;
    my $self = {
	workers => [],
	maxworkers => 2,
	pending => [],
	handler => sub { die "This should not happen"; },
    };
    foreach my $p (qw(maxworkers handler))
    {
	$self->{$p} = $params{$p} if $params{$p};
    }
    return bless $self, $class;
}

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

    while (scalar @{$self->{workers}} < $self->{maxworkers})
    {
	my $w = Cassandane::Unit::Worker->new();
	$w->{handler} = $self->{handler};
	$w->start();
	push(@{$self->{workers}}, $w);
    }
}

# Assign an work item to an idle worker if necessary
# block until a worker is idle.
sub assign
{
    my ($self, $witem) = @_;

    my @idle = grep { !$_->{busy}; } @{$self->{workers}};
    my $w = shift @idle || $self->_wait();
    $w->assign($witem);
}

# Wait for a Worker to send back a completed work item.
# Mark the Worker idle, remember its work item where
# retrieve() will find it, and returns the Worker.
sub _wait
{
    my ($self) = @_;


    # Build the bit mask for select()
    my $rbits = '';
    foreach my $w (@{$self->{workers}})
    {
	next if (!$w->{busy});
	vec($rbits, fileno($w->{uppipe}), 1) = 1;
    }

    # select() with no timeout
    my $res;
    do {
	$res = select($rbits, undef, undef, undef);
    } while ($res < 0 && $! == EINTR);
    die "select failed: $!" if ($res < 0);

    # discover which of our workers has responded
    foreach my $w (@{$self->{workers}})
    {
	if (vec($rbits, fileno($w->{uppipe}), 1))
	{
	    push(@{$self->{pending}}, $w->get_reply());
	    return $w;
	}
    }
    die "Unexpected result from select: $res";
}

# Retrieve a completed work item.  If $blocking is true,
# wait if necessary (used when draining i.e. no more work
# items will be made available).
sub retrieve
{
    my ($self, $blocking) = @_;

    if ($blocking && !scalar @{$self->{pending}})
    {
	my @busy = grep { $_->{busy}; } @{$self->{workers}};
	$self->_wait() if (scalar @busy);
    }
    return shift @{$self->{pending}};
}

# reap all workers
sub stop
{
    my ($self) = @_;

    while (my $w = pop @{$self->{workers}})
    {
	$w->stop();
    }
}

sub DESTROY
{
    my ($self) = @_;
    $self->stop();
}

package Cassandane::Unit::WorkerListener;
use base qw(Test::Unit::Listener);
use Cassandane::Util::Log;

sub new
{
    my ($class) = shift;
    my $self = {
	witem => undef,
    };
    return bless $self, $class;
}

sub start_suite
{
    my ($self, $suite) = @_;
    # nothing to see here
}

sub start_test
{
    my ($self, $test) = @_;
    # nothing to see here
}

sub end_test
{
    my ($self, $test) = @_;
    # nothing to see here
}

sub end_suite
{
    my ($self, $suite) = @_;
    # nothing to see here
}

sub add_error
{
    my ($self, $test, $exception) = @_;
    my $witem = $self->{witem};
    $witem->{result} = 'error';

    # Remove '-object' which points at the TestCase, which will have all
    # sorts of stuff we can't thaw.  We have enough information to
    # discover the right TestCase in the parent process.
    $exception->{'-object'} = undef;
    $witem->{exception} = $exception;
}

sub add_failure
{
    my ($self, $test, $exception) = @_;
    my $witem = $self->{witem};
    $witem->{result} = 'fail';

    # Remove '-object' which points at the TestCase, which will have all
    # sorts of stuff we can't thaw.  We have enough information to
    # discover the right TestCase in the parent process.
    $exception->{'-object'} = undef;
    $witem->{exception} = $exception;
}

sub add_pass
{
    my ($self, $test) = @_;
    my $witem = $self->{witem};
    $witem->{result} = 'pass';
}

package Cassandane::Unit::TestPlan;
use File::Temp qw(tempfile);
use File::Path qw(mkpath);
use Cassandane::Util::Log;

my @test_roots = (
    'Cassandane/Test',
    'Cassandane/Cyrus',
);

sub new
{
    my ($class, %opts) = @_;
    my $self = {
	schedule => {},
	keep_going => delete $opts{keep_going} || 0,
	log_directory => delete $opts{log_directory},
	maxworkers => delete $opts{maxworkers} || 1,
    };
    die "Unknown options: " . join(' ', keys %opts)
	if scalar %opts;
    return bless $self, $class;
}

sub _get_item
{
    my ($self, $suite) = @_;
    return $self->{schedule}->{$suite} ||=
	Cassandane::Unit::TestPlanItem->new($suite);
}

sub _schedule
{
    my ($self, $neg, $path, $testname) = @_;
    return if ($path =~ m/\/TestCase\.pm$/);

    my $suite = $path;
    $suite =~ s/\.pm$//;
    $suite =~ s/\//::/g;

    if ($neg eq '!')
    {
	if (defined $testname)
	{
	    # disable a specific test
	    $self->_get_item($suite)->_deny($testname);
	}
	else
	{
	    # remove entire suite
	    delete $self->{schedule}->{$suite};
	}
    }
    else
    {
	# add to the schedule
	my $item = $self->_get_item($suite);
	if (defined $testname)
	{
	    $item->_allow($testname) if $testname;
	}
    }
}

# Returns ($neg, $ostype, $ospath, $testname)
sub _parse_test_spec
{
    my ($name) = @_;

    my ($neg, $path) = ($name =~ m/^([~!]?)(.*)$/);
    $path =~ s/\.pm$//g;
    $path =~ s/::/\//g;
    $path =~ s/\./\//g;
    $path =~ s/\/+/\//g;
    $path =~ s/^\/*//;
    $path =~ s/\/*$//;

    $neg = '!' if $neg eq '~';

    foreach my $root (@test_roots)
    {
	return ($neg, 'd', $path, undef)
	    if ($root eq $path);

	my $fpath = $path;
	$fpath = "$root/$path"
	    if ("$root/" ne substr($path, 0, length($root)+1));

	return ($neg, 'd', $fpath, undef)
	    if ( -d $fpath );
	return ($neg, 'f', "$fpath.pm", undef)
	    if ( -f "$fpath.pm" );

	my $test;
	($fpath, $test) = ($fpath =~ m/^(.*)\/([^\/]+)$/);
	next unless defined $test;
	return ($neg, 'f', "$fpath.pm", $test)
	    if ( -f "$fpath.pm" );
    }

    die "Unrecognised test specification: $name";
}

sub schedule
{
    my ($self, @names) = @_;

    @names = @test_roots
	if !scalar @names;

    foreach my $name (@names)
    {
	my ($neg, $type, $path, $test) = _parse_test_spec($name);

	$self->schedule(@test_roots)
	    if $neg eq '!' && !scalar %{$self->{schedule}};

	if ($type eq 'd')
	{
	    opendir DIR, $path
		or die "Cannot open directory $path for reading: $!";
	    while ($_ = readdir DIR)
	    {
		next unless m/\.pm$/;
		$self->_schedule($neg, "$path/$_", undef);
	    }
	    closedir DIR;
	}
	else
	{
	    $self->_schedule($neg, $path, $test);
	}
    }
}


#
# Get the entire expanded schedule as specific {suite,testname} tuples,
# sorted in alphabetic order on suite name then testname.
#
sub _get_schedule
{
    my ($self) = @_;

    my @items = sort { $a->{suite} cmp $b->{suite} } values %{$self->{schedule}};
    my @res;
    foreach my $item (@items)
    {
	my $loaded = $item->_get_loaded_suite();
	foreach my $name (sort @{$loaded->names()})
	{
	    $name =~ s/^test_//;
	    next unless $item->_is_allowed($name);

	    my (@settings) = Cassandane::Unit::TestCase->make_parameter_settings($item->{suite});
	    foreach my $setting (@settings)
	    {
		push(@res, {
		    suite => $item->{suite},
		    testname => $name,
		    result => 'unknown',
		    exception => undef,
		    logfile => undef,
		    parameter_setting => $setting,
		});
	    }
	}
    }
    return @res;
}

# Sort and return the schedule as a list of "suite.test" strings
# e.g. "Cassandane::Cyrus::Quota.using_storage".
sub list
{
    my ($self) = @_;

    my @res;
    foreach my $witem ($self->_get_schedule())
    {
	push(@res, "$witem->{suite}.$witem->{testname}");
    }

    return @res;
}

sub _setup_logfile
{
    my ($self, $witem) = @_;

    # Flush the old stdout/stderr
    ${\*STDOUT}->flush;
    ${\*STDERR}->flush;

    # Save the old stdout/stderr - this is important
    # for the single threaded case where inter-test
    # messages go to the original stdout.
    open my $oldout, '>&', \*STDOUT
	or die "Cannot save STDOUT";
    open my $olderr, '>&', \*STDERR
	or die "Cannot save STDERR";

    my $logfh;
    my $logfile;
    srand;   # XXX does this fix the tempfile() issue (#42)?
    if (defined $self->{log_directory})
    {
	# Log directory specified so create the log file
	# there with a semi-obvious name
	if (! -d $self->{log_directory})
	{
	    mkpath($self->{log_directory})
		or die "Cannot create directory $self->{log_directory}: $!";
	}
	my $template = $witem->{suite} .  '.' .  $witem->{testname} .  '.XXXXXX';
	$template =~ s/::/./g;
	($logfh, $logfile) = tempfile($template,
				      DIR => $self->{log_directory},
				      SUFFIX => '.log',
				      UNLINK => 0);
	chmod(0644, $logfile);
    }
    else
    {
	# Create a per-test temporary logfile
	($logfh, $logfile) = tempfile(UNLINK => 0);
    }

    # Redirect both STDOUT and STDERR to the log file
    open STDOUT, '>&', $logfh
	or die "Cannot redirect STDOUT";
    open STDERR, '>&', $logfh
	or die "Cannot redirect STDERR";
    close $logfh;

    $witem->{logfile} = $logfile;
    $self->{oldout} = $oldout;
    $self->{olderr} = $olderr;
}

# Redirect STDOUT and STDERR back to their original fds
sub _restore_stdout
{
    my ($self) = @_;

    ${\*STDOUT}->flush;
    open STDOUT, '>&', $self->{oldout}
	or die "Cannot restore STDOUT";
    close $self->{oldout};
    $self->{oldout} = undef;

    ${\*STDERR}->flush;
    open STDERR, '>&', $self->{olderr}
	or die "Cannot restore STDERR";
    close $self->{olderr};
    $self->{olderr} = undef;
}

sub _dump_logfile
{
    my ($logfile) = @_;

    open LOGFILE, '<', $logfile
	or die "Cannot open $logfile for reading: $!";
    while (<LOGFILE>)
    {
	print STDERR $_;
    }
    close LOGFILE;
}

sub _get_suite_and_test
{
    my ($self, $witem) = @_;
    my $suite = $self->_get_item($witem->{suite})->_get_loaded_suite();
    my ($test) = grep { $_->name() eq 'test_' . $witem->{testname}; } @{$suite->tests()};
    return ($suite, $test);
}

sub _run_workitem
{
    my ($self, $witem, $result, $runner, $annotate_flag) = @_;
    my ($suite, $test) = $self->_get_suite_and_test($witem);
    Cassandane::Unit::TestCase->enable_test($witem->{testname});
    $self->_setup_logfile($witem);
    Cassandane::Unit::TestCase->apply_parameter_setting($witem->{parameter_setting});
    $suite->run($result, $runner);

    if ($test->can('post_tear_down'))
    {
	eval
	{
	    $test->post_tear_down();
	};
	my $ex = $@;
	if ($ex)
	{
	    $result->add_error($test,
			       Test::Unit::Error->make_new_from_error($ex));
	}
    }

    $self->_restore_stdout();
    if ($annotate_flag)
    {
	$test->annotate_from_file($witem->{logfile});
	_dump_logfile($witem->{logfile}) if (get_verbose > 1);
	unlink($witem->{logfile}) if (!defined $self->{log_directory});
    }
}

sub _finish_workitem
{
    my ($self, $witem, $result, $runner) = @_;
    my ($suite, $test) = $self->_get_suite_and_test($witem);

    $result->start_test($test);
    if ($runner->can('fake_start_time'))
    {
	$runner->fake_start_time($test, $witem->{start_time});
    }

    $test->annotate_from_file($witem->{logfile});
    _dump_logfile($witem->{logfile}) if (get_verbose > 1);
    unlink($witem->{logfile}) if (!defined $self->{log_directory});

    if ($witem->{result} eq 'pass')
    {
	$result->add_pass($test);
    }
    elsif ($witem->{result} eq 'fail')
    {
	$witem->{exception}->{'-object'} = $test;
	$result->add_failure($test, $witem->{exception});
    }
    elsif ($witem->{result} eq 'error')
    {
	$witem->{exception}->{'-object'} = $test;
	$result->add_error($test, $witem->{exception});
    }
    $result->end_test($test);
}

sub _setup_worker_listeners
{
    my ($result, $wlistener) = @_;

    # Remove the output format listener
    my @list;
    my $found = 0;
    foreach my $ll (@{$result->{_Listeners}})
    {
	push(@list, $ll)
	    unless defined $ll->{remove_me_in_cassandane_child};
	$found ||= (ref($ll) eq ref($wlistener));
    }
    push(@list, $wlistener)
	unless $found;
    $result->{_Listeners} = \@list;
}

# The 'run' method makes this class look sufficiently like a
# Test::Unit::TestCase that Test::Unit::TestRunner will happily run it.
# This enables us to run all our scheduled tests with a single
# TestResult and a single summary of errors.
sub run
{
    my ($self, $result, $runner) = @_;

    my $maxworkers = $self->{maxworkers} || 1;

    # we expand the schedule before forking the
    # workers so that we can just hand the reference
    # to the worker
    my @workitems = $self->_get_schedule();

    # try to clean up after ourselves on interrupt
    my $interrupted = 0;
    $SIG{INT} = sub {
	$interrupted ++;
	# third ^C will terminate without cleanup
	$SIG{INT} = 'DEFAULT' if $interrupted >= 2;
    };

    if ($maxworkers > 1)
    {
	# multi-threaded case: use worker pool

	# we want an error not a signal
	$SIG{PIPE} = 'IGNORE';

	# Just In Case any code samples this in a TestCase c'tor
	$ENV{TEST_UNIT_WORKER_ID} = 'invalid';

	my $wlistener = Cassandane::Unit::WorkerListener->new();

	my $pool = Cassandane::Unit::WorkerPool->new(
	    maxworkers => $maxworkers,
	    handler => sub {
		my ($witem) = @_;
		$wlistener->{witem} = $witem;
		_setup_worker_listeners($result, $wlistener);
		$self->_run_workitem($witem, $result, $runner, 0);
	    },
	);
	my $witem;
	$pool->start();
	# first ^C stops spawning new work items
	while ($interrupted < 1 && ($witem = shift @workitems))
	{
	    $pool->assign($witem)
		if ($self->{keep_going} || $result->was_successful());
	    while ($witem = $pool->retrieve(0))
	    {
		$self->_finish_workitem($witem, $result, $runner);
	    }
	}
	# second ^C stops waiting for work items to finish
	while ($interrupted < 2 && ($witem = $pool->retrieve(1)))
	{
	    $self->_finish_workitem($witem, $result, $runner);
	}
	$pool->stop();
    }
    else
    {
	# single threaded case: just run it all in-process
	foreach my $witem (@workitems)
	{
	    $self->_run_workitem($witem, $result, $runner, 1);
	    last if ($interrupted || !($self->{keep_going} || $result->was_successful()));
	}
    }

    return $result->was_successful();
}

1;
