import example

assert int(example.square(2)) == 4
assert int(example.square(3)) == 9