Some small issues fix

And clean up some test code