Add fixes and tests