From 5da37422878179483c72e49c4e4575ed0146c052 Mon Sep 17 00:00:00 2001
From: Aristo Chen <aristo.chen@canonical.com>
Date: Wed, 17 Jul 2024 10:11:41 +0800
Subject: [PATCH 1/2] Support provision tool(genio-flash) for MediaTek devices

* Add baudrate 921600 for MediaTek devices
* Disable Pylint R0913 to allow more args in the function. By
  default Pylint only allows 5 argument in the function
* Add a new key(extra_provision_tool_args) for user to specify
  additional arguments consumed by provision tool. For example
  we can use something like "-e list_dtbo=aaaa.dtbo bbbb.dtbo"
  with MediaTek devices to make it load different dtbo when
  doing test.

Signed-off-by: Aristo Chen <aristo.chen@canonical.com>
---
 README                    |  2 +-
 sanity/agent/agent.py     |  1 +
 sanity/agent/deploy.py    | 51 +++++++++++++++++++++++++++++++++++++--
 sanity/launcher/parser.py |  6 ++++-
 4 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/README b/README
index 2fa09d3..95f98f4 100644
--- a/README
+++ b/README
@@ -12,7 +12,7 @@ Program your own tplan yaml
             password: target device user password
             serial_console:
                 port: /dev/ttyUSBxxx
-                baud_rate: baud rate, support[115200, 9600]
+                baud_rate: baud rate, support[115200, 9600, 921600]
             network: target device interface
             extra-recepients: ["yourmail@mail.com"]
         
diff --git a/sanity/agent/agent.py b/sanity/agent/agent.py
index 79b3cc3..16ff8e3 100644
--- a/sanity/agent/agent.py
+++ b/sanity/agent/agent.py
@@ -75,6 +75,7 @@ def start(plan, con, sched=None):
                         con,
                         stage["deploy"].get("utility"),
                         stage["deploy"].get("method"),
+                        stage["deploy"].get("extra_provision_tool_args"),
                         stage["deploy"].get("update_boot_assets", False),
                         stage["deploy"].get("timeout", 600),
                     )
diff --git a/sanity/agent/deploy.py b/sanity/agent/deploy.py
index a2f9d18..600aa68 100644
--- a/sanity/agent/deploy.py
+++ b/sanity/agent/deploy.py
@@ -166,11 +166,58 @@ def boot_assets_update(addr):
     syscmd(cmd)
 
 
-# pylint: disable=R0911,R0912,R0914,R0915
-def deploy(con, method, user_init, update_boot_assets, timeout=600):
+# pylint: disable=R0911,R0912,R0913,R0914,R0915
+def deploy(
+    con,
+    method,
+    user_init,
+    extra_provision_tool_args,
+    update_boot_assets,
+    timeout=600,
+):
     """handle different provision method and run provision"""
     files = "seed"
     match method:
+        case "genio_flash":
+            # Assuming the host is >= 22.04
+            syscmd(
+                "set -x; sudo apt-get install -y python3-pip python3-pip-whl"
+            )
+            gitlab_url = (
+                "git+https://gitlab.com/mediatek/aiot/bsp/"
+                + "genio-tools.git#egg=genio-tools"
+            )
+            syscmd(f"set -x; pip3 install -U -e {gitlab_url}")
+            syscmd("set -x; export PATH=$PATH:/home/$USER/.local/bin/")
+            syscmd("set -x; genio-config")
+            image_tarball = [
+                f
+                for f in os.listdir("./")
+                if re.search(
+                    r"genio-(core|classic-(server|desktop)).*\.(tar\.xz)$",
+                    f,
+                )
+            ][0]
+            image_dir = image_tarball.split(".")[0]
+            boot_assets_tarball = [
+                f
+                for f in os.listdir("./")
+                if re.search(
+                    r"genio.*boot-assets.*\.(tar\.xz)$",
+                    f,
+                )
+            ][0]
+            syscmd(f"set -x; rm -rf {image_dir}")
+            syscmd(f"set -x; tar xf {image_tarball}")
+            syscmd(
+                "set -x; tar --strip-components=1 -xf "
+                f"{boot_assets_tarball} -C {image_dir}"
+            )
+            syscmd(
+                f"set -x; cd {image_dir};"
+                f"genio-flash '{extra_provision_tool_args}'"
+            )
+
         case "utp_com":
             # This method is currently used for i.MX6 devices that does not
             # use uuu to flash image
diff --git a/sanity/launcher/parser.py b/sanity/launcher/parser.py
index 15a6bab..11fa429 100644
--- a/sanity/launcher/parser.py
+++ b/sanity/launcher/parser.py
@@ -21,7 +21,7 @@
                         "port": {"type": "string"},
                         "baud_rate": {
                             "type": "integer",
-                            "enum": [115200, 9600],
+                            "enum": [115200, 9600, 921600],
                             "default": 115200,
                         },
                     },
@@ -60,6 +60,7 @@
                                     "utility": {
                                         "type": "string",
                                         "enum": [
+                                            "genio_flash",
                                             "utp_com",
                                             "uuu",
                                             "uuu_bootloader",
@@ -69,6 +70,9 @@
                                         ],
                                     },
                                     "method": {"$ref": "#/$defs/method"},
+                                    "extra_provision_tool_args": {
+                                        "type": "string"
+                                    },
                                     "timeout": {
                                         "type": "integer",
                                         "default": 600,

From e39c52f29e84095ddc626057e91a38ff4a6897de Mon Sep 17 00:00:00 2001
From: Aristo Chen <aristo.chen@canonical.com>
Date: Thu, 5 Sep 2024 11:20:01 +0800
Subject: [PATCH 2/2] Move whl into pyproject.toml

Signed-off-by: Aristo Chen <aristo.chen@canonical.com>
---
 pyproject.toml         | 1 +
 sanity/agent/deploy.py | 3 ---
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index cae2ce0..f250c5c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -8,6 +8,7 @@ dependencies = [
     "pyserial",
     "schedule",
     "wrapt_timeout_decorator",
+    "whl",
 ]
 
 dynamic = ["version"]
diff --git a/sanity/agent/deploy.py b/sanity/agent/deploy.py
index 600aa68..a899d86 100644
--- a/sanity/agent/deploy.py
+++ b/sanity/agent/deploy.py
@@ -180,9 +180,6 @@ def deploy(
     match method:
         case "genio_flash":
             # Assuming the host is >= 22.04
-            syscmd(
-                "set -x; sudo apt-get install -y python3-pip python3-pip-whl"
-            )
             gitlab_url = (
                 "git+https://gitlab.com/mediatek/aiot/bsp/"
                 + "genio-tools.git#egg=genio-tools"