Working on testing
Test on models pass