Code changes + adding tests