Cleaning up code to satisfy flake