[Feature]: We need to add some unit tests

## Motivation

Having unit tests can help us make sure that the code is still working while we are making updates.

## Proposed solution

Consult with Codex to see what unit tests can be added.

## Alternatives considered

Not applied
