Minor changes and updates to locomotive unit test