I would think this would be the expected behaviour would be this instead:
www.example.com -> example.com
www.example.com/test.html -> example.com/test.html
That's a good idea. Just map the last parts of the URL (particularly folder and file) from the old domain to the new domain. You can easily use mod-rewrite if you have apache. Just make an .htaccess file in the folder where the www.example.com document root is and add the following:
RewriteEngine On
RewriteCond %{HTTP_HOST} ^www\.example\.com$ [NC]
RewriteRule ^(.*)$ http://example.com/$1 [R=301,L]
That way, everything will be nicely redirected over from the www version to the non-www version.
My only concern is, is it correct for it to use a 301 redirect to bring it back to the front-page no matter where it has come from?
The only time you should redirect a URL to the homepage is if that URL was the homepage in the past. For example, if your previous homepage was at the url http://www.example.com/homepage.html and the new homepage is just example.com, then you can use a 301 return code on http://www.example.com/homepage.html and redirect it to example.com.
Make sure you test everything as you go because you don't want to end up with too many redirects or you'll see a message in your browser something similar to "the document has moved here" and when you click on "here", the same message reappears.