Write more tests and modify User model