Custom AndroidJUnitRunner with wake lock

CustomJUnitRunner.java

import static android.view.WindowManager.LayoutParams.FLAG_DISMISS_KEYGUARD;
import static android.view.WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON;
import static android.view.WindowManager.LayoutParams.FLAG_TURN_SCREEN_ON;

public class CustomJUnitRunner extends AndroidJUnitRunner {

  @Override public void onCreate(Bundle arguments) {
    super.onCreate(arguments);

    ActivityLifecycleMonitorRegistry.getInstance().addLifecycleCallback(new ActivityLifecycleCallback() {
      @Override public void onActivityLifecycleChanged(Activity activity, Stage stage) {
        if (stage == Stage.PRE_ON_CREATE) {
          activity.getWindow().addFlags(FLAG_DISMISS_KEYGUARD | FLAG_TURN_SCREEN_ON | FLAG_KEEP_SCREEN_ON);
        }
      }
    });
  }
}

Hello,
Thank you for the code.
Can you tell me please how can I use it in Android Studio in a uiautomator test? or where shoudl I add this code?
Thank you for the help.
My best regards,
Mostafa