Added further functionality to our test