Update tests to work with total code rewrite