Kea2-python 0.1.3__py3-none-any.whl → 0.2.1__py3-none-any.whl

kea2/fastbotManager.py

@@ -1,13 +1,17 @@
+from retry import retry
+from retry.api import retry_call
 from dataclasses import asdict
-import subprocess
-import threading
 import requests
 from time import sleep
-from .adbUtils import push_file
+
+
+from uiautomator2.core import HTTPResponse, _http_request
+from kea2.adbUtils import ADBDevice, ADBStreamShell_V2
 from pathlib import Path
-from .utils import getLogger
+from kea2.utils import getLogger
+
 
-from typing import IO, TYPE_CHECKING
+from typing import IO, TYPE_CHECKING, Dict
 if TYPE_CHECKING:
     from .keaUtils import Options
 
@@ -21,56 +25,49 @@ class FastbotManager:
         self.log_file: str = log_file
         self.port = None
         self.thread = None
+        self._device_output_dir = None
+        ADBDevice.setDevice(options.serial, options.transport_id)
+        self.dev = ADBDevice()
 
-
-    def _activateFastbot(self) -> threading.Thread:
+    def _activateFastbot(self) -> ADBStreamShell_V2:
         """
         activate fastbot.
         :params: options: the running setting for fastbot
         :params: port: the listening port for script driver
         :return: the fastbot daemon thread
         """
-        options = self.options
         cur_dir = Path(__file__).parent
-        push_file(
+        self.dev.sync.push(
             Path.joinpath(cur_dir, "assets/monkeyq.jar"),
-            "/sdcard/monkeyq.jar",
-            device=options.serial
+            "/sdcard/monkeyq.jar"
         )
-        push_file(
+        self.dev.sync.push(
             Path.joinpath(cur_dir, "assets/fastbot-thirdpart.jar"),
             "/sdcard/fastbot-thirdpart.jar",
-            device=options.serial,
         )
-        push_file(
+        self.dev.sync.push(
             Path.joinpath(cur_dir, "assets/kea2-thirdpart.jar"),
             "/sdcard/kea2-thirdpart.jar",
-            device=options.serial,
         )
-        push_file(
-            Path.joinpath(cur_dir, "assets/framework.jar"), 
+        self.dev.sync.push(
+            Path.joinpath(cur_dir, "assets/framework.jar"),
             "/sdcard/framework.jar",
-            device=options.serial
         )
-        push_file(
-            Path.joinpath(cur_dir, "assets/fastbot_libs/arm64-v8a"),
-            "/data/local/tmp",
-            device=options.serial
+        self.dev.sync.push(
+            Path.joinpath(cur_dir, "assets/fastbot_libs/arm64-v8a/libfastbot_native.so"),
+            "/data/local/tmp/arm64-v8a/libfastbot_native.so",
         )
-        push_file(
-            Path.joinpath(cur_dir, "assets/fastbot_libs/armeabi-v7a"),
-            "/data/local/tmp",
-            device=options.serial
+        self.dev.sync.push(
+            Path.joinpath(cur_dir, "assets/fastbot_libs/armeabi-v7a/libfastbot_native.so"),
+            "/data/local/tmp/armeabi-v7a/libfastbot_native.so",
         )
-        push_file(
-            Path.joinpath(cur_dir, "assets/fastbot_libs/x86"),
-            "/data/local/tmp",
-            device=options.serial
+        self.dev.sync.push(
+            Path.joinpath(cur_dir, "assets/fastbot_libs/x86/libfastbot_native.so"),
+            "/data/local/tmp/x86/libfastbot_native.so",
         )
-        push_file(
-            Path.joinpath(cur_dir, "assets/fastbot_libs/x86_64"),
-            "/data/local/tmp",
-            device=options.serial
+        self.dev.sync.push(
+            Path.joinpath(cur_dir, "assets/fastbot_libs/x86_64/libfastbot_native.so"),
+            "/data/local/tmp/x86_64/libfastbot_native.so",
         )
 
         t = self._startFastbotService()
@@ -79,29 +76,84 @@ class FastbotManager:
         return t
 
 
-    def check_alive(self, port):
+    def check_alive(self):
         """
         check if the script driver and proxy server are alive.
         """
-        for _ in range(10):
-            sleep(2)
-            try:
-                requests.get(f"http://localhost:{port}/ping")
-                return
-            except requests.ConnectionError:
-                logger.info("waiting for connection.")
-                pass
-        raise RuntimeError("Failed to connect fastbot")
-
-
-    def _startFastbotService(self) -> threading.Thread:
+        def _check_alive_request():
+            _http_request(dev=self.dev, device_port=8090, method="GET", path="/ping")
+
+        try:
+            retry_call(_check_alive_request, tries=10, delay=2)
+        except requests.ConnectionError:
+            raise RuntimeError("Failed to connect fastbot")
+
+    def request(self, method: str, path: str, data: Dict=None, timeout: int=10) -> HTTPResponse:
+        return _http_request(self.dev, 8090, method, path, data, timeout)
+
+    @retry(Exception, tries=2, delay=2)
+    def init(self, options: "Options", stamp):
+        post_data = {
+            "takeScreenshots": options.take_screenshots,
+            "Stamp": stamp,
+            "deviceOutputRoot": options.device_output_root,
+        }
+        r = _http_request(
+            self.dev,
+            device_port=8090,
+            method="POST",
+            path="/init",
+            data=post_data
+        )
+        print(f"[INFO] Init fastbot: {post_data}", flush=True)
+        import re
+        self._device_output_dir = re.match(r"outputDir:(.+)", r.text).group(1)
+        print(f"[INFO] Fastbot initiated. outputDir: {r.text}", flush=True)
+    
+    @retry(Exception, tries=2, delay=2)
+    def stepMonkey(self, monkeyStepInfo):
+        r = self.request(
+            method="POST",
+            path="/stepMonkey",
+            data=monkeyStepInfo
+        )
+        return r.json()["result"]
+
+    @retry(Exception, tries=2, delay=2)
+    def stopMonkey(self):
+        """
+        send a stop monkey request to the server.
+        """
+        r = self.request(
+            method="GET",
+            path="/stopMonkey",
+        )
+
+        print(f"[Server INFO] {r.text}", flush=True)
+    
+    @retry(Exception, tries=2, delay=2)
+    def logScript(self, execution_info: Dict):
+        r = self.request(
+            method="POST",
+            path="/logScript",
+            data=execution_info
+        )
+        res = r.text
+        if res != "OK":
+            print(f"[ERROR] Error when logging script: {execution_info}", flush=True)
+
+    @property
+    def device_output_dir(self):
+        return self._device_output_dir
+
+    def _startFastbotService(self) -> ADBStreamShell_V2:
         shell_command = [
             "CLASSPATH="
             "/sdcard/monkeyq.jar:"
             "/sdcard/framework.jar:"
             "/sdcard/fastbot-thirdpart.jar:"
             "/sdcard/kea2-thirdpart.jar",
-            
+
             "exec", "app_process",
             "/system/bin", "com.android.commands.monkey.Monkey",
             "-p", *self.options.packageNames,
@@ -119,37 +171,36 @@ class FastbotManager:
 
         full_cmd = ["adb"] + (["-s", self.options.serial] if self.options.serial else []) + ["shell"] + shell_command
 
+
         outfile = open(self.log_file, "w", encoding="utf-8", buffering=1)
 
         logger.info("Options info: {}".format(asdict(self.options)))
         logger.info("Launching fastbot with shell command:\n{}".format(" ".join(full_cmd)))
         logger.info("Fastbot log will be saved to {}".format(outfile.name))
 
-        # process handler
-        proc = subprocess.Popen(full_cmd, stdout=outfile, stderr=outfile)
-        t = threading.Thread(target=self.close_on_exit, args=(proc, outfile), daemon=True)
-        t.start()
-
+        t = self.dev.stream_shell(shell_command, stdout=outfile, stderr=outfile)
         return t
 
-    def close_on_exit(self, proc: subprocess.Popen, f: IO):
+    def close_on_exit(self, proc: ADBStreamShell_V2, f: IO):
         self.return_code = proc.wait()
         f.close()
         if self.return_code != 0:
             raise RuntimeError(f"Fastbot Error: Terminated with [code {self.return_code}] See {self.log_file} for details.")
 
     def get_return_code(self):
-        if self.thread:
+        if self.thread.is_running():
             logger.info("Waiting for Fastbot to exit.")
-            self.thread.join()
-        return self.return_code
+            return self.thread.wait()
+        return self.thread.poll()
 
     def start(self):
+        # kill the fastbot process if runing.
+        self.dev.kill_proc("com.android.commands.monkey")
         self.thread = self._activateFastbot()
 
     def join(self):
-        if self.thread:
-            self.thread.join()
+        self.thread.join()
+
 
 
 

kea2/keaUtils.py

@@ -2,6 +2,7 @@ import json
 import os
 from pathlib import Path
 import traceback
+import time
 from typing import Callable, Any, Dict, List, Literal, NewType, TypedDict, Union
 from unittest import TextTestRunner, registerResult, TestSuite, TestCase, TextTestResult
 import random
@@ -110,6 +111,8 @@ class Options:
     packageNames: List[str]
     # target device
     serial: str = None
+    # target device with transport_id
+    transport_id: str = None
     # test agent. "native" for stage 1 and "u2" for stage 1~3
     agent: Literal["u2", "native"] = "u2"
     # max step in exploration (availble in stage 2~3)
@@ -139,8 +142,13 @@ class Options:
     def __post_init__(self):
         import logging
         logging.basicConfig(level=logging.DEBUG if self.debug else logging.INFO)
-        if self.serial and self.Driver:
-            self.Driver.setDeviceSerial(self.serial)
+        if self.Driver:
+            target_device = dict()
+            if self.serial:
+                target_device["serial"] = self.serial
+            if self.transport_id:
+                target_device["transport_id"] = self.transport_id
+            self.Driver.setDevice(target_device)
         global LOGFILE, RESFILE, STAMP
         if self.log_stamp:
             illegal_chars = ['/', '\\', ':', '*', '?', '"', '<', '>', '|', '\n', '\r', '\t', '\0']
@@ -313,10 +321,11 @@ class KeaTestRunner(TextTestRunner):
                 result.flushResult(outfile=RESFILE)
                 # setUp for the u2 driver
                 self.scriptDriver = self.options.Driver.getScriptDriver()
-                fb.check_alive(port=self.scriptDriver.lport)
-                self._init()
+                fb.check_alive()
+                
+                fb.init(options=self.options, stamp=STAMP)
 
-                resultSyncer = ResultSyncer(self.device_output_dir, self.options.output_dir)
+                resultSyncer = ResultSyncer(fb.device_output_dir, self.options)
                 resultSyncer.run()
 
                 end_by_remote = False
@@ -331,7 +340,7 @@ class KeaTestRunner(TextTestRunner):
                     )
 
                     try:
-                        xml_raw = self.stepMonkey()
+                        xml_raw = fb.stepMonkey(self._monkeyStepInfo)
                         propsSatisfiedPrecond = self.getValidProperties(xml_raw, result)
                     except requests.ConnectionError:
                         logger.info("Connection refused by remote.")
@@ -369,18 +378,18 @@ class KeaTestRunner(TextTestRunner):
                     print("execute property %s." % execPropName, flush=True)
 
                     result.addExcuted(test)
-                    self._logScript(result.lastExecutedInfo)
+                    fb.logScript(result.lastExecutedInfo)
                     try:
                         test(result)
                     finally:
                         result.printErrors()
 
                     result.updateExectedInfo()
-                    self._logScript(result.lastExecutedInfo)
+                    fb.logScript(result.lastExecutedInfo)
                     result.flushResult(outfile=RESFILE)
 
                 if not end_by_remote:
-                    self.stopMonkey()
+                    fb.stopMonkey()
                 result.flushResult(outfile=RESFILE)
                 resultSyncer.close()
                 
@@ -424,40 +433,46 @@ class KeaTestRunner(TextTestRunner):
         self.stream.flush()
         return result
 
-    def stepMonkey(self) -> str:
-        """
-        send a step monkey request to the server and get the xml string.
-        """
+    # def stepMonkey(self) -> str:
+    #     """
+    #     send a step monkey request to the server and get the xml string.
+    #     """
+    #     block_dict = self._getBlockedWidgets()
+    #     block_widgets: List[str] = block_dict['widgets']
+    #     block_trees: List[str] = block_dict['trees']
+    #     URL = f"http://localhost:{self.scriptDriver.lport}/stepMonkey"
+    #     logger.debug(f"Sending request: {URL}")
+    #     logger.debug(f"Blocking widgets: {block_widgets}")
+    #     logger.debug(f"Blocking trees: {block_trees}")
+    #     r = requests.post(
+    #         url=URL,
+    #         json={
+    #             "steps_count": self.stepsCount,
+    #             "block_widgets": block_widgets,
+    #             "block_trees": block_trees
+    #         }
+    #     )
+
+    #     res = json.loads(r.content)
+    #     xml_raw = res["result"]
+    #     return xml_raw
+
+    @property
+    def _monkeyStepInfo(self):
+        r = self._get_block_widgets()
+        r["steps_count"] = self.stepsCount
+        return r
+    
+    def _get_block_widgets(self):
         block_dict = self._getBlockedWidgets()
         block_widgets: List[str] = block_dict['widgets']
         block_trees: List[str] = block_dict['trees']
-        URL = f"http://localhost:{self.scriptDriver.lport}/stepMonkey"
-        logger.debug(f"Sending request: {URL}")
         logger.debug(f"Blocking widgets: {block_widgets}")
         logger.debug(f"Blocking trees: {block_trees}")
-        r = requests.post(
-            url=URL,
-            json={
-                "steps_count": self.stepsCount,
-                "block_widgets": block_widgets,
-                "block_trees": block_trees
-            }
-        )
-
-        res = json.loads(r.content)
-        xml_raw = res["result"]
-        return xml_raw
-
-    def stopMonkey(self) -> str:
-        """
-        send a stop monkey request to the server and get the xml string.
-        """
-        URL = f"http://localhost:{self.scriptDriver.lport}/stopMonkey"
-        logger.debug(f"Sending request: {URL}")
-        r = requests.get(URL)
-
-        res = r.content.decode(encoding="utf-8")
-        print(f"[Server INFO] {res}", flush=True)
+        return {
+            "block_widgets": block_widgets,
+            "block_trees": block_trees
+        }
 
     def getValidProperties(self, xml_raw: str, result: JsonResult) -> PropertyStore:
 
@@ -490,40 +505,13 @@ class KeaTestRunner(TextTestRunner):
                     print(f"{getFullPropName(test)} has reached its max_tries. Skip.", flush=True)
                     continue
                 validProps[propName] = test
-        
+
         print(f"{len(validProps)} precond satisfied.", flush=True)
         if len(validProps) > 0:
             print("[INFO] Valid properties:",flush=True)
             print("\n".join([f'                - {getFullPropName(p)}' for p in validProps.values()]), flush=True)
         return validProps
 
-    def _logScript(self, execution_info:Dict):
-        URL = f"http://localhost:{self.scriptDriver.lport}/logScript"
-        r = requests.post(
-            url=URL,
-            json=execution_info
-        )
-        res = r.content.decode(encoding="utf-8")
-        if res != "OK":
-            print(f"[ERROR] Error when logging script: {execution_info}", flush=True)
-
-    def _init(self):
-        URL = f"http://localhost:{self.scriptDriver.lport}/init"
-        data = {
-            "takeScreenshots": self.options.take_screenshots,
-            "Stamp": STAMP,
-            "deviceOutputRoot": self.options.device_output_root,
-        }
-        print(f"[INFO] Init fastbot: {data}", flush=True)
-        r = requests.post(
-            url=URL,
-            json=data
-        )
-        res = r.content.decode(encoding="utf-8")
-        import re
-        self.device_output_dir = re.match(r"outputDir:(.+)", res).group(1)
-        print(f"[INFO] Fastbot initiated. outputDir: {res}", flush=True)
-
     def collectAllProperties(self, test: TestSuite):
         """collect all the properties to prepare for PBT
         """
@@ -679,17 +667,19 @@ class KeaTestRunner(TextTestRunner):
     def __del__(self):
         """tearDown method. Cleanup the env.
         """
-        try:
-            self.stopMonkey()
-        except Exception as e:
-            pass
-
         if self.options.Driver:
             self.options.Driver.tearDown()
 
         try:
+            start_time = time.time()
             logger.info("Generating bug report")
             report_generator = BugReportGenerator(self.options.output_dir)
             report_generator.generate_report()
+
+            end_time = time.time()
+            generation_time = end_time - start_time
+
+            logger.info(f"Bug report generation completed in {generation_time:.2f} seconds")
+            
         except Exception as e:
             logger.error(f"Error generating bug report: {e}", flush=True)

kea2/kea_launcher.py

@@ -9,10 +9,22 @@ def _set_runner_parser(subparsers: "argparse._SubParsersAction[argparse.Argument
         "-s",
         "--serial",
         dest="serial",
+        required=False,
+        default=None,
         type=str,
         help="The serial of your device. Can be found with `adb devices`",
     )
 
+    parser.add_argument(
+        "-t",
+        "--transport-id",
+        dest="transport_id",
+        required=False,
+        default=None,
+        type=str,
+        help="transport-id of your device, can be found with `adb devices -l`",
+    )
+
     parser.add_argument(
         "-p",
         "--packages",
@@ -128,6 +140,8 @@ def driver_info_logger(args):
     print("[INFO] Driver Settings:", flush=True)
     if args.serial:
         print("  serial:", args.serial, flush=True)
+    if args.transport_id:
+        print("  transport_id:", args.transport_id, flush=True)
     if args.package_names:
         print("  package_names:", args.package_names, flush=True)
     if args.agent:
@@ -174,6 +188,7 @@ def run(args=None):
         Driver=U2Driver,
         packageNames=args.package_names,
         serial=args.serial,
+        transport_id=args.transport_id,
         running_mins=args.running_minutes,
         maxStep=args.max_step,
         throttle=args.throttle_ms,

kea2/logWatcher.py

@@ -2,7 +2,8 @@ import re
 import os
 import threading
 import time
-from .utils import getLogger
+from typing import IO
+from kea2.utils import getLogger
 
 
 logger = getLogger(__name__)
@@ -20,47 +21,43 @@ def thread_excepthook(args):
 
 class LogWatcher:
 
-    def watcher(self, poll_interval=0.5):
-        self.buffer = ""
+    def watcher(self, poll_interval=3):
         self.last_pos = 0
 
-        while not self.end_flag:
-            self.read_log()
-            time.sleep(poll_interval)
+        with open(self.log_file, "r") as fp:
+            while not self.end_flag:
+                self.read_log(fp)
+                time.sleep(poll_interval)
+            
+            time.sleep(0.2)
+            self.read_log(fp)
         
-        time.sleep(0.2)
-        self.read_log()
-        
-    def read_log(self):
-        with open(self.log_file, 'r', encoding='utf-8') as f:
-            f.seek(self.last_pos)
-            new_data = f.read()
-            self.last_pos = f.tell()
-
-            if new_data:
-                self.buffer += new_data
-            self.parse_log()
-
-    def parse_log(self):
-        buffer = self.buffer
-        exception_match = PATTERN_EXCEPTION.search(buffer)
+    def read_log(self, f: IO):
+        f.seek(self.last_pos)
+        buffer = f.read()
+        self.last_pos = f.tell()
+
+        self.parse_log(buffer)
+
+    def parse_log(self, content):
+        exception_match = PATTERN_EXCEPTION.search(content)
         if exception_match:
             exception_body = exception_match.group(1).strip()
             if exception_body:
                 raise RuntimeError(
-                    "[Error] Execption while running fastbot:\n" + 
+                    "[Error] Fatal Execption while running fastbot:\n" + 
                     exception_body + 
                     "\nSee fastbot.log for details."
                 )
-        if self.end_flag:
-            statistic_match = PATTERN_STATISTIC.search(buffer)
-            if statistic_match:
-                statistic_body = statistic_match.group(1).strip()
-                if statistic_body:
-                    print(
-                        "[INFO] Fastbot exit:\n" + 
-                        statistic_body
-                    , flush=True)
+        
+        statistic_match = PATTERN_STATISTIC.search(content)
+        if statistic_match:
+            statistic_body = statistic_match.group(1).strip()
+            if statistic_body:
+                print(
+                    "[INFO] Fastbot exit:\n" + 
+                    statistic_body
+                , flush=True)
 
     def __init__(self, log_file):
         logger.info(f"Watching log: {log_file}")
@@ -79,4 +76,4 @@ class LogWatcher:
 
 
 if __name__ == "__main__":
-    LogWatcher("fastbot.log")
+    LogWatcher("/Users/atria/Desktop/coding/Kea2/output/res_2025062510_0420056539/fastbot_2025062510_0420056539.log")

kea2/resultSyncer.py

@@ -1,19 +1,26 @@
+from pathlib import Path
 import threading
-from .adbUtils import adb_shell, pull_file
+from .adbUtils import ADBDevice
 from .utils import getLogger
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from .keaUtils import Options
 
 logger = getLogger(__name__)
 
 
 class ResultSyncer:
 
-    def __init__(self, device_output_dir, output_dir):
+    def __init__(self, device_output_dir, options: "Options"):
         self.device_output_dir = device_output_dir
-        self.output_dir = output_dir
+        self.output_dir = Path(options.output_dir) / Path(device_output_dir).name
         self.running = False
         self.thread = None
         self.sync_event = threading.Event()
 
+        ADBDevice.setDevice(serial=options.serial, transport_id=options.transport_id)
+        self.dev = ADBDevice()
+
     def run(self):
         """Start a background thread to sync device data when triggered"""
         self.running = True
@@ -37,7 +44,8 @@ class ResultSyncer:
         try:
             logger.debug(f"Removing device output directory: {self.device_output_dir}")
             remove_device_dir = ["rm", "-rf", self.device_output_dir]
-            adb_shell(remove_device_dir)
+            # adb_shell(remove_device_dir)
+            self.dev.shell(remove_device_dir)
         except Exception as e:
             logger.error(f"Error removing device output directory: {e}", flush=True)
 
@@ -48,9 +56,11 @@ class ResultSyncer:
         try:
             logger.debug("Syncing data")
 
-            pull_file(self.device_output_dir, str(self.output_dir))
+            self.dev.sync.pull_dir(self.device_output_dir, self.output_dir, exist_ok=True)
+            # pull_file(self.device_output_dir, str(self.output_dir))
 
-            remove_pulled_screenshots = ["find", self.device_output_dir, "-name", "\"*.png\"", "-delete"]
-            adb_shell(remove_pulled_screenshots)
+            remove_pulled_screenshots = ["find", self.device_output_dir, "-name", '"*.png"', "-delete"]
+            self.dev.shell(remove_pulled_screenshots)
+            # adb_shell(remove_pulled_screenshots)
         except Exception as e:
-            logger.error(f"Error in data sync: {e}", flush=True)
+            logger.error(f"Error in data sync: {e}")