def square(x):
  return x * x

def test_square():
  assert square(2) == 4

def test_square_negative():
  assert square(-2) == 4