Added tests & modified unit generator.