Using Mockolate to proxy calls from a mocked Socket to a ByteArray

Alecmce_Sockets.as

package mockolate.issues {
	import mockolate.*;
	import mockolate.ingredients.*;
	import mockolate.runner.*;
	import flash.net.Socket;
	import flash.utils.ByteArray;
	import org.hamcrest.*;
	import org.hamcrest.core.*;
	import org.hamcrest.object.*;

	public class Alecmce_Sockets {

		[Rule]
		public var mocks:MockolateRule = new MockolateRule();

		[Mock]
		public var socket:Socket;

		public var bytes:ByteArray;

		[Test]
		public function is_not_explody():void {
			assertThat(socket, not(nullValue()));
		}

		[Test]
		public function readByte_should_be_proxied():void {
			bytes = new ByteArray();
			bytes.writeByte(8);
			bytes.position = 0;

			proxy(socket.readByte, bytes.readByte);

			assertThat(socket.readByte(), equalTo(8));
		}

		[Test]
		public function writeByte_should_be_proxied():void {
			bytes = new ByteArray();

			proxy(socket.writeByte, bytes.writeByte, [arg(int)]);

			socket.writeByte(11);

			assertThat(bytes.position, equalTo(1));
			bytes.position = 0;
			assertThat(bytes.readByte(), equalTo(11));
		}

		// Proxy very specific methods. Must supply matchers for all arguments. 
		// 
		private function proxy(from:Function, to:Function, args:Array = null):void {
			expect(from.apply(null, args))
				.callsWithInvocation(function(invocation:Invocation):* {
					return to.apply(null, invocation.arguments);
				});
		}

		[Test]
		public function reads_should_be_proxied():void {
			bytes = new ByteArray();
			bytes.writeByte(101);
			bytes.writeInt(303);
			bytes.writeUTF("SH2000");
			bytes.position = 0;

			proxyAny(socket, bytes, ['readByte', 'readInt', 'readUTF']);

			assertThat(socket.readByte(), equalTo(101));
			assertThat(socket.readInt(), equalTo(303));
			assertThat(socket.readUTF(), equalTo("SH2000"));
		}

		// Proxy arbitrary methods
		// 
		private function proxyAny(source:*, target:*, methods:Array):void {
			methods.forEach(function(method:String, i:int, a:Array):void {
				mock(source)
					.method(method)
					.anyArgs()
					.callsWithInvocation(function(invocation:Invocation):* {
						return target[method].apply(target, invocation.arguments);
					});
			});
		}
	}
}

That's pretty sexy.