refactor: init uitest service

Author: codematrixer

README.md

@@ -200,10 +200,13 @@ d.uninstall_app("com.kuaishou.hmapp")
 ### 启动App
 
 ```python
+d.start_app("com.kuaishou.hmapp")
+
 d.start_app("com.kuaishou.hmapp", "EntryAbility")
 ```
-传入的两个参数分别是`package_name`, `page_name`，可以通过hdc命令获取`hdc shell aa dump -l`
+`package_name`, `page_name`分别为包名和ability name，可以通过hdc命令获取：`hdc shell aa dump -l`
 
+不传`page_name`时，默认会使用main ability作为`page_name`
 
 ### 停止App
 ```python
@@ -258,6 +261,23 @@ d.get_app_info("com.kuaishou.hmapp")
 }
 ```
 
+### 获取App main ability
+```python
+d.get_app_main_ability("com.kuaishou.hmapp")
+```
+
+输出的数据结构是Dict, 内容如下
+
+```
+{
+    "name": "EntryAbility",
+    "moduleName": "kwai",
+    "moduleMainAbility": "EntryAbility",
+    "mainModule": "kwai",
+    "isLauncherAbility": true,
+    "score": 2
+}
+```
 
 ## 设备操作
 ### 获取设备信息

hmdriver2/_client.py

@@ -1,12 +1,11 @@
 # -*- coding: utf-8 -*-
-import re
 import socket
 import json
 import time
 import os
-import glob
+import hashlib
 import typing
-import subprocess
+from typing import Optional
 from datetime import datetime
 from functools import cached_property
 
@@ -143,8 +142,7 @@ def invoke_captures(self, api: str, args: typing.List = []) -> HypiumResponse:
 
     def start(self):
         logger.info("Start HmClient connection")
-        self._init_so_resource()
-        self._restart_uitest_service()
+        _UITestService(self.hdc).init()
 
         self._connect_sock()
 
@@ -163,47 +161,79 @@ def release(self):
             logger.error(f"An error occurred: {e}")
 
     def _create_hdriver(self) -> DriverData:
-        logger.debug("create uitest driver")
+        logger.debug("Create uitest driver")
         resp: HypiumResponse = self.invoke("Driver.create")  # {"result":"Driver#0"}
         hdriver: DriverData = DriverData(resp.result)
         return hdriver
 
-    def _init_so_resource(self):
-        "Initialize the agent.so resource on the device."
-
-        file_postfix = ".so"
-        device_agent_path = "/data/local/tmp/agent.so"
-        arch_info = self.hdc.shell("file /system/bin/uitest").output.strip()
-        if "x86_64" in arch_info:
-            file_postfix = ".x86_64_so"
-        local_path = ""
-        local_ver = "0.0.0"
-        for agent_file in glob.glob(os.path.join(ASSETS_PATH, "uitest_agent*so")):
-            file_name = os.path.split(agent_file)[1]
-            if not agent_file.endswith(file_postfix):
-                continue
-            matcher = re.search(r'\d{1,3}[.]\d{1,3}[.]\d{1,3}', file_name)
-            if not matcher:
-                continue
-            ver = matcher.group()[0]
-            if ver.split('.') > local_ver.split('.'):
-                local_ver, local_path = ver, agent_file
-        device_ver_info = self.hdc.shell(f"cat {device_agent_path} | grep -a UITEST_AGENT_LIBRARY").output.strip()
-        matcher = re.search(r'\d{1,3}[.]\d{1,3}[.]\d{1,3}', device_ver_info)
-        device_ver = matcher.group(0) if matcher else "0.0.0"
-        logger.debug(f"local agent version {local_ver}, device agent version {device_ver}")
-        if device_ver.split('.') < local_ver.split('.'):
-            logger.debug(f"start update agent, path is {local_path}")
-            self._kill_uitest_service()
-            for file in AGENT_CLEAR_PATH:
-                self.hdc.shell(f"rm /data/local/tmp/{file}*")
-            self.hdc.send_file(local_path, device_agent_path)
-            self.hdc.shell(f"chmod +x {device_agent_path}")
-            logger.debug("Update agent finish.")
-        else:
-            logger.debug("Device agent is up to date!")
-
-    def get_devicetest_proc_pid(self):
+
+class _UITestService:
+    def __init__(self, hdc: HdcWrapper):
+        """Initialize the UITestService class."""
+        self.hdc = hdc
+
+    def init(self):
+        """
+        Initialize the UITest service:
+        1. Ensure agent.so is set up on the device.
+        2. Start the UITest daemon.
+
+        Note: 'hdc shell aa test' will also start a uitest daemon.
+        $ hdc shell ps -ef |grep uitest
+        shell        44306     1 25 11:03:37 ?    00:00:16 uitest start-daemon singleness
+        shell        44416     1 2 11:03:42 ?     00:00:01 uitest start-daemon com.hmtest.uitest@4x9@1"
+        """
+
+        logger.debug("Initializing UITest service")
+        local_path = self._get_local_agent_path()
+        remote_path = "/data/local/tmp/agent.so"
+
+        self._kill_uitest_service()  # Stop the service if running
+        self._setup_device_agent(local_path, remote_path)
+        self._start_uitest_daemon()
+        time.sleep(0.5)
+
+    def _get_local_agent_path(self) -> str:
+        """Return the local path of the agent file."""
+        target_agent = "uitest_agent_v1.1.0.so"
+        return os.path.join(os.path.dirname(os.path.realpath(__file__)), "assets", target_agent)
+
+    def _get_remote_md5sum(self, file_path: str) -> Optional[str]:
+        """Get the MD5 checksum of a remote file."""
+        command = f"md5sum {file_path}"
+        output = self.hdc.shell(command).output.strip()
+        return output.split()[0] if output else None
+
+    def _get_local_md5sum(self, file_path: str) -> str:
+        """Get the MD5 checksum of a local file."""
+        hash_md5 = hashlib.md5()
+        with open(file_path, "rb") as f:
+            for chunk in iter(lambda: f.read(4096), b""):
+                hash_md5.update(chunk)
+        return hash_md5.hexdigest()
+
+    def _is_remote_file_exists(self, file_path: str) -> bool:
+        """Check if a file exists on the device."""
+        command = f"[ -f {file_path} ] && echo 'exists' || echo 'not exists'"
+        result = self.hdc.shell(command).output.strip()
+        return "exists" in result
+
+    def _setup_device_agent(self, local_path: str, remote_path: str):
+        """Ensure the remote agent file is correctly set up."""
+        if self._is_remote_file_exists(remote_path):
+            local_md5 = self._get_local_md5sum(local_path)
+            remote_md5 = self._get_remote_md5sum(remote_path)
+            if local_md5 == remote_md5:
+                logger.debug("Remote agent file is up-to-date")
+                self.hdc.shell(f"chmod +x {remote_path}")
+                return
+            self.hdc.shell(f"rm {remote_path}")
+
+        self.hdc.send_file(local_path, remote_path)
+        self.hdc.shell(f"chmod +x {remote_path}")
+        logger.debug("Updated remote agent file")
+
+    def _get_uitest_pid(self) -> typing.List[str]:
         proc_pids = []
         result = self.hdc.shell("ps -ef").output.strip()
         lines = result.splitlines()
@@ -215,23 +245,11 @@ def get_devicetest_proc_pid(self):
         return proc_pids
 
     def _kill_uitest_service(self):
-        for pid in self.get_devicetest_proc_pid():
+        for pid in self._get_uitest_pid():
             self.hdc.shell(f"kill -9 {pid}")
             logger.debug(f"Killed uitest process with PID {pid}")
 
-    def _restart_uitest_service(self):
-        """
-        Restart the UITest daemon.
-
-        Note: 'hdc shell aa test' will also start a uitest daemon.
-        $ hdc shell ps -ef |grep uitest
-        shell        44306     1 25 11:03:37 ?    00:00:16 uitest start-daemon singleness
-        shell        44416     1 2 11:03:42 ?     00:00:01 uitest start-daemon com.hmtest.uitest@4x9@1"
-        """
-        try:
-            self._kill_uitest_service()
-            self.hdc.shell("uitest start-daemon singleness")
-            time.sleep(.5)
-
-        except subprocess.CalledProcessError as e:
-            logger.error(f"An error occurred: {e}")
+    def _start_uitest_daemon(self):
+        """Start the UITest daemon."""
+        self.hdc.shell("uitest start-daemon singleness")
+        logger.debug("Started UITest daemon")

hmdriver2/assets/uitest_agent_v1.0.8.x86_64_so

@@ -83,11 +83,20 @@ def _invoke(self, api: str, args: List = []) -> HypiumResponse:
         return self._client.invoke(api, this="Driver#0", args=args)
 
     @delay
-    def start_app(self, package_name: str, page_name: str = ''):
-        page_name = page_name or self.get_app_main_ability(package_name).get('name', 'MainAbility')
+    def start_app(self, package_name: str, page_name: Optional[str] = None):
+        """
+        Start an application on the device.
+        If the `package_name` is empty, it will retrieve main ability using `get_app_main_ability`.
+
+        Args:
+            package_name (str): The package name of the application.
+            page_name (Optional[str]): Ability Name within the application to start.
+        """
+        if not page_name:
+            page_name = self.get_app_main_ability(package_name).get('name', 'MainAbility')
         self.hdc.start_app(package_name, page_name)
 
-    def force_start_app(self, package_name: str, page_name: str = ""):
+    def force_start_app(self, package_name: str, page_name: Optional[str] = None):
         self.go_home()
         self.stop_app(package_name)
         self.start_app(package_name, page_name)
@@ -153,9 +162,8 @@ def get_app_abilities(self, package_name: str) -> List[Dict]:
         """
         Get the abilities of an application.
 
-
         Args:
-            package_name (str): The package name of the application to retrieve information for.
+            package_name (str): The package name of the application.
 
         Returns:
             List[Dict]: A list of dictionaries containing the abilities of the application.

hmdriver2/driver.py

@@ -114,6 +114,7 @@ def recv_file(self, rpath: str, lpath: str):
         return result
 
     def shell(self, cmd: str, error_raise=True) -> CommandResult:
+        # ensure the command is wrapped in double quotes
         if cmd[0] != '\"':
             cmd = "\"" + cmd
         if cmd[-1] != '\"':

hmdriver2/hdc.py

@@ -2,3 +2,4 @@
 
 pytest tests --capture=no
 # pytest tests/test_driver.py::test_toast --capture=no
+# pytest tests/test_driver.py::test_toast tests/test_driver.py::test_get_app_main_ability --capture=no

runtest.sh

@@ -32,6 +32,12 @@ def test_force_start_app(d):
     d.force_start_app("com.samples.test.uitest", "EntryAbility")
 
 
+def test_get_app_main_ability(d):
+    d.unlock()
+    ability = d.get_app_main_ability("com.samples.test.uitest")
+    assert ability.get("name") == "EntryAbility"
+
+
 def test_clear_app(d):
     d.clear_app("com.samples.test.uitest")