Fixed a bug and added visual testing