Write unit tests

[vishnu-deepsource]

Currently, the project lacks proper unit tests. There are some end-to-end tests, however they are not easy to set up. We should add proper unit tests for the project.