Added tests

Tests were forgotten in previous commit