No public descriptionmultitest-transport-dev

PiperOrigin-RevId: 628358040
Change-Id: I11583d8af6315feb823c7702963968e03257534d

2 files changed, 43 insertions, 18 deletions

diff --git a/multitest_transport/cli/cli.py b/multitest_transport/cli/cli.py
index 5587ee8..8b32cc6 100644
--- a/multitest_transport/cli/cli.py
+++ b/multitest_transport/cli/cli.py
@@ -98,6 +98,8 @@ _ALLOW_TO_UPDATE_KEY = 'allowToUpdate'
 # string in startConsole() method after console.start();
 # in tools/tradefederation/core/src/com/android/tradefed/command/Console.java
 _TF_CONSOLE_SUCCESS_INDICATOR = 'help all'
+# Success indicator once omni lab server started.
+_OMNI_LAB_SERVER_SUCCESS_INDICATOR = 'Lab server successfully started'
 # command for check log: "docker logs mtt"
 _DOCKER_LOGS_MTT_COMMAND = ['logs', 'mtt']
 # interval in second for checking out the logs
@@ -367,34 +369,42 @@ def _CheckDockerImageVersion(docker_helper, container_name):
         '(%s < %s)' % (cli_version, image_version))
 
 
-def _IsTfConsoleSuccessfullyStarted(host):
-  """Check is TF console started successfully.
+def _IsConsoleSuccessfullyStarted(host, is_omnilab_based):
+  """Check is the success indicator detected from docker logs.
 
   Args:
     host: an instance of host_util.Host.
+    is_omnilab_based: is omnilab based console or not.
+
   Returns:
-    True if find the TF console success indicator in docker logs.
+    True if find the success indicator in docker logs.
   Raises:
     RuntimeError: if exceptions detected in docker logs.
   """
+  indicator = (
+      _OMNI_LAB_SERVER_SUCCESS_INDICATOR
+      if is_omnilab_based
+      else _TF_CONSOLE_SUCCESS_INDICATOR
+  )
   docker_context = command_util.DockerContext(host.context, login=False)
   end_time = time.time() + _MTT_SERVER_WAIT_TIME_SECONDS
   while time.time() <= end_time:
     remaining_time = int(end_time - time.time())
     docker_context.RequestTfConsolePrintOut()
-    command_result = docker_context.Run(_DOCKER_LOGS_MTT_COMMAND,
-                                        timeout=remaining_time)
+    command_result = docker_context.Run(
+        _DOCKER_LOGS_MTT_COMMAND, timeout=remaining_time
+    )
     docker_log = command_result.stderr + '\n' + command_result.stdout
 
-    if _TF_CONSOLE_SUCCESS_INDICATOR in command_result.stdout:
+    if indicator in command_result.stdout:
       return True
     elif 'exception' in docker_log.lower():
-      raise RuntimeError(
-          f'Tradefed failed to start with exception:\n{docker_log}')
+      raise RuntimeError(f'ATS failed to start with exception:\n{docker_log}')
     time.sleep(_LOG_INQUIRE_INTERVAL_SEC)
   # when timeout
   raise RuntimeError(
-      'ATS replica failed to start in %ss' % _MTT_SERVER_WAIT_TIME_SECONDS)
+      'ATS replica failed to start in %ss' % _MTT_SERVER_WAIT_TIME_SECONDS
+  )
 
 
 def Start(args, host=None):
@@ -675,7 +685,7 @@ def _StartMttNode(args, host):
     # localhost URL.
     hostname = 'localhost'
   if control_server_url:
-    if _IsTfConsoleSuccessfullyStarted(host):
+    if _IsConsoleSuccessfullyStarted(host, args.is_omnilab_based):
       logger.info('ATS replica is running.')
   else:
     url = 'http://%s:%s' % (hostname, args.port)
diff --git a/multitest_transport/cli/cli_test.py b/multitest_transport/cli/cli_test.py
index ce51a57..20c6f15 100644
--- a/multitest_transport/cli/cli_test.py
+++ b/multitest_transport/cli/cli_test.py
@@ -96,7 +96,8 @@ class CliTest(parameterized.TestCase):
         return_value=self.mock_control_server_client)
     self.submit_host_update_event_patcher.start()
     self.mock_tf_console_started_patcher = mock.patch.object(
-        cli, '_IsTfConsoleSuccessfullyStarted', return_value=True)
+        cli, '_IsConsoleSuccessfullyStarted', return_value=True
+    )
     self.mock_tf_console_started_patcher.start()
     self.mock_tf_console_print_out_patcher = mock.patch.object(
         command_util.DockerContext, 'RequestTfConsolePrintOut',
@@ -241,27 +242,41 @@ class CliTest(parameterized.TestCase):
     cli._CheckDockerImageVersion(docker_helper, container_name)
 
   @mock.patch.object(command_util.DockerContext, 'Run')
-  def testCheckTfConsoleSuccessfullyStarted_withSuccessIndicator(
+  def testCheckConsoleSuccessfullyStarted_withSuccessIndicator(
       self, mock_command_result):
     self.mock_tf_console_started_patcher.stop()
     mock_command_result.return_value = common.CommandResult(
         return_code=0, stdout=cli._TF_CONSOLE_SUCCESS_INDICATOR, stderr='')
-    self.assertTrue(cli._IsTfConsoleSuccessfullyStarted(self._CreateHost()))
+    self.assertTrue(
+        cli._IsConsoleSuccessfullyStarted(self._CreateHost(), False)
+    )
+
+  @mock.patch.object(command_util.DockerContext, 'Run')
+  def testCheckConsoleSuccessfullyStarted_withOmniBasedSuccessIndicator(
+      self, mock_command_result
+  ):
+    self.mock_tf_console_started_patcher.stop()
+    mock_command_result.return_value = common.CommandResult(
+        return_code=0, stdout=cli._OMNI_LAB_SERVER_SUCCESS_INDICATOR, stderr=''
+    )
+    self.assertTrue(
+        cli._IsConsoleSuccessfullyStarted(self._CreateHost(), True)
+    )
 
   @mock.patch.object(command_util.DockerContext, 'Run')
-  def testCheckTfConsoleSuccessfullyStarted_withException(
+  def testCheckConsoleSuccessfullyStarted_withException(
       self, mock_command_result):
     self.mock_tf_console_started_patcher.stop()
     stderr = 'com.android.tradefed.config.ConfigurationException'
     mock_command_result.return_value = common.CommandResult(
         return_code=0, stdout='', stderr=stderr)
     with self.assertRaisesRegex(
-        RuntimeError, r'.*Tradefed failed to start with exception:*'):
-      cli._IsTfConsoleSuccessfullyStarted(self._CreateHost())
+        RuntimeError, r'.*ATS failed to start with exception:*'):
+      cli._IsConsoleSuccessfullyStarted(self._CreateHost(), False)
 
   @mock.patch.object(time, 'time')
   @mock.patch.object(command_util.DockerContext, 'Run')
-  def testCheckTfConsoleSuccessfullyStarted_withTimeout(
+  def testCheckConsoleSuccessfullyStarted_withTimeout(
       self, mock_command_result, mock_time):
     self.mock_tf_console_started_patcher.stop()
     mock_command_result.return_value = common.CommandResult(
@@ -269,7 +284,7 @@ class CliTest(parameterized.TestCase):
     mock_time.side_effect = iter([0, cli._MTT_SERVER_WAIT_TIME_SECONDS + 1])
     with self.assertRaisesRegex(
         RuntimeError, r'.*ATS replica failed to start in*'):
-      cli._IsTfConsoleSuccessfullyStarted(self._CreateHost())
+      cli._IsConsoleSuccessfullyStarted(self._CreateHost(), False)
 
   def testStart(self):
     """Test start without service account."""