Add more tests

[mostafa]

Ticket(s)

Closes #430

Description

Add more tests.

Related PRs

N/A

Development Checklist

• I have added a descriptive title to this PR.
• I have squashed related commits together.
• I have rebased my branch on top of the latest main branch.
• I have performed a self-review of my own code.
• I have commented on my code, particularly in hard-to-understand areas.
• I have added docstring(s) to my code.
• I have made corresponding changes to the documentation (docs).
• I have updated docs using make gen-docs command.
• I have added tests for my changes.
• I have signed all the commits.

Legal Checklist

• I have read the CONTRIBUTING document.
• I have read and understood the Code of Conduct.
• I have read and agreed to the Apache CLA (required).
• I have read and agreed to the LICENSE (required).