added some more code, in theory works, but need to write tests now
adding unittest