fix: unit tests

[athith-g]

Description

Fixes #214.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• My code follows the style guidelines of this project
• I have performed a self-review of my own code
• I have updated the documentation (or documentation does not need to be updated)
• My changes generate no new warnings
• New and existing unit tests pass locally with my changes
• I have not reduced the existing code coverage

[sourcery-ai]

Hey @athith-g - I've reviewed your changes and they look great!

Here's what I looked at during the review

• 🟢 General issues: all looks good
• 🟢 Security: all looks good
• 🟢 Testing: all looks good
• 🟢 Complexity: all looks good

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment to tell me if it was helpful.}