Updated with more tests