Fixing tests and updating reviewer feedback.