**Is your feature request related to a problem? Please describe.** <!-- A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] --> Right now the code coverage is very low and we should improve it by adding unit tests wherever possible.