added model tests, improved code based on tests