The implementation for this is awkward. The junit step does indeed set the stage status to UNSTABLE, but it also quickly sets the whole build status to UNSTABLE.
Indeed, there is not really any easy way to truly fix the issue, so the reality is that there are multiple kinds of "results" in Pipelines: Run.getResult, ErrorAction, and now WarningAction.
Can the junit plugin only change the status for the stage it is called in?
The main reason why we couldn't do something like that is that it would be a breaking change for everyone that currently checks things like currentBuild.currentResult inside of the Pipeline. Step-level "results" handled via WarningAction have no relation to that result, and there is no easy way to fix that, because we need a FlowNode context to get a WarningAction, which cannot be provided by currentBuild.currentResult.
When changing the whole build status to UNSTABLE, the fact that it also changed the stage status is not visible anyway.
As far as I know, in Blue Ocean and the Pipeline Steps view, the overall build result no longer changes the displayed result for individual steps and stages.
If you are talking about Pipeline Stage View, then yes, that plugin needs to be updated to use WarningAction correctly, that is tracked as JENKINS-58783.