Fixing test.py for Squares.