Fixing None return in Controller