Made changes as suggested by @Ocramius
Tests passing