Cleaning up stream stuff and added test