chore(core): unit tests: use `km_core_state_context_set_if_needed` 🌱

If we have unit tests that use legacy `set_context` functions then we should move them to `km_core_state_context_set_if_needed` so that we can remove the legacy APIs.