Add more unit tests to improve code coverage

You can check codecov runs for the lastest master branch and find lines of code that aren't currently covered by any unit test. It would be very helpful if you can add a unit test to cover such lines of code.

Important: any new unit tests you add must run very quickly! Ideally in under 1sec. If your tests run on some data, try to use the smallest dataset possible for the test to still be useful.