Cleaned up the code to pass more network tests