cheets_CTS/GTS: Keep screen on for more "Media" tests.

Previous implementation focused only on modules that
uses additional media asset downloads as an approximation,
but this isn't enough given the expanded test set in R.

BUG=b:175569272
TEST=cheets_CTS_R.11_r3.arm.CtsMediaV2TestCases

Change-Id: I9d8b90f5b5aee68a7f8f22b56739e21d96db08eb
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/autotest/+/2591838
Commit-Queue: Kazuhiro Inaba <kinaba@chromium.org>
Commit-Queue: Shao-Chuan Lee <shaochuan@chromium.org>
Tested-by: Kazuhiro Inaba <kinaba@chromium.org>
Auto-Submit: Kazuhiro Inaba <kinaba@chromium.org>
Reviewed-by: Shao-Chuan Lee <shaochuan@chromium.org>
diff --git a/server/cros/tradefed/tradefed_test.py b/server/cros/tradefed/tradefed_test.py
index e76e463..f7e45b1 100644
--- a/server/cros/tradefed/tradefed_test.py
+++ b/server/cros/tradefed/tradefed_test.py
@@ -1261,23 +1261,21 @@
                     command = self._tradefed_retry_command(retry_template,
                                                            session_id)
 
-                # TODO(pwang): Evaluate if it is worth it to get the number of
-                #              not-excecuted, for instance, by collecting all
-                #              tests on startup (very expensive, may take 30
-                #              minutes).
                 if media_asset and media_asset.uri:
                     # Clean-up crash logs from previous sessions to ensure
                     # enough disk space for 16GB storage devices: b/156075084.
                     if not keep_media:
                         self._clean_crash_logs()
-                    # TODO(b/137917339): Only prevent screen from turning off for
-                    # media tests. Remove this check once the GPU issue is fixed.
+                # TODO(b/137917339): Only prevent screen from turning off for
+                # media tests. Remove this check once the GPU issue is fixed.
+                keep_screen_on = (media_asset and media_asset.uri) or (
+                        target_module and "Media" in target_module)
+                if keep_screen_on:
                     self._override_powerd_prefs()
                 try:
                     waived_tests, acc = self._run_and_parse_tradefed(command)
                 finally:
-                    # TODO(b/137917339): ditto
-                    if media_asset and media_asset.uri:
+                    if keep_screen_on:
                         self._restore_powerd_prefs()
                 if media_asset:
                     self._fail_on_unexpected_media_download(media_asset)