Improve test flakiness

What does this MR do and why?

Fix feature category and improve test flakiness by replacing multiple find_by with a single query.

Changelog: changed
EE: true

References

Screenshots or screen recordings

Before	After

How to set up and validate locally

MR acceptance checklist

Evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.

Related to #551198 (closed)