PR changes, fix tests and new tests