Added assert statements to test the function.