Rebasing and fixing tests on master.
Adding changes discussed in #358