AsyncRule.as

package utils {
	
import org.flexunit.internals.runners.statements.IAsyncStatement;
import org.flexunit.internals.runners.statements.MethodRuleBase;
import org.flexunit.rules.IMethodRule;
import org.flexunit.runners.model.FrameworkMethod;
import org.flexunit.token.AsyncTestToken;
import org.flexunit.token.ChildResult;
	

public class AsyncRule extends MethodRuleBase implements IMethodRule {

	protected var hasWrappedFunctions:Boolean;
	protected var isComplete:Boolean;

	public function wrap(closure:Function):Function {
		hasWrappedFunctions = true;

		function asyncRuleWrapHelper(...args):* {
			var result:ChildResult;
			var returnValue:*;

			try {
				returnValue = closure.apply(null, args);
			} catch (e:Error) {
				complete();
				sendComplete(e);
				return undefined;
			}

			if (isComplete) {
				sendComplete(null);
				return undefined;
			}

			return returnValue;
		}
		
		return asyncRuleWrapHelper;
	}

	public function complete():void {
		if (!hasWrappedFunctions)
			throw Error("Calling 'complete' without wrapping any functions?");

		// Once we are 'complete', the termination of the next 'wrapped'
		// function will finish the test.
		isComplete = true;
	}
	
	override public function evaluate(parentToken:AsyncTestToken):void {
		super.evaluate(parentToken);
		
		hasWrappedFunctions = false;
		isComplete = false;
		
		proceedToNextStatement();
	}

	override protected function handleStatementComplete(result:ChildResult):void {
		if (result.error) {
			// If we got an error in the test method, fail immediately.
			if (hasWrappedFunctions)
				complete();
			sendComplete(result.error);
		}

		// If we didn't wrap any functions, we're done.
		if (!hasWrappedFunctions || isComplete)
			sendComplete(result.error);
		
		// Otherwise, do nothing until we get a call to 'complete'.
	}
}
}

TestAsyncRule.as

package tests {
import flash.errors.IOError;
import flash.utils.setTimeout;

import utils.AsyncRule;
import utils.assertEqual;

public class TestAsyncRule {
	
	[Rule]
	public var async:AsyncRule = new AsyncRule();
	
	[Test]
	public function testWrapPassthroughArgs():void {
		function toWrap(arg:*):* {
			return arg;
		}
		
		var wrapped:Function = async.wrap(toWrap);
		assertEqual(wrapped("hello"), "hello");
		async.complete();
	}
	
	[Test(expects="flash.errors.IOError")]
	public function testWrapErrors():void {
		setTimeout(async.wrap(function():void {
			throw new IOError("some error");
		}), 100);
	}
	
}
}