Also, I added some tests for the new functionality of my previous commit