Add a simple unittest instead of the previous integration test