Fix tests and bug in learn code