Updated the testing logic

[levtelyatnikov]

Checklist

• My pull request has a clear and explanatory title.
• My pull request passes the Linting test.
• I added appropriate unit tests and I made sure the code passes all unit tests. (refer to comment below)
• My PR follows PEP8 guidelines. (refer to comment below)
• My code is properly documented, using numpy docs conventions, and I made sure the documentation renders properly.
• I linked to issues and PRs that are relevant to this PR.

Description

Issue

Additional context

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.01%. Comparing base (21c9340) to head (aba78b2).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #255   +/-   ##
=======================================
  Coverage   94.01%   94.01%           
=======================================
  Files         184      184           
  Lines        6664     6664           
=======================================
  Hits         6265     6265           
  Misses        399      399           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.