esp-test-utils 0.2.0__tar.gz → 0.2.1__tar.gz

{esp_test_utils-0.2.0 → esp_test_utils-0.2.1}/.gitignore

@@ -14,10 +14,13 @@ gl-codequality.json
 /dist/
 *.egg-info/
 
+# tests
 /log/
 /logs/
 dut_logs/
 dut.log
+tests/utility/_files/test-bin
+
 
 # docs
 /docs/*.rst

{esp_test_utils-0.2.0 → esp_test_utils-0.2.1}/.pre-commit-config.yaml

@@ -35,6 +35,7 @@ repos:
       - 'types-psutil'
       - 'types-PyYAML'
       - 'typing_extensions'
+      - 'sqlalchemy'
 
   - repo: https://github.com/espressif/conventional-precommit-linter
     rev: v1.10.0
@@ -48,3 +49,9 @@ repos:
       - id: codespell
         args: ["--write-changes"]
         additional_dependencies: [tomli]
+
+exclude: |
+  (
+    ^docs/.*
+    | gen_esp32part\.py$
+  )

{esp_test_utils-0.2.0 → esp_test_utils-0.2.1}/CHANGELOG.md

@@ -1,3 +1,12 @@
+## v0.2.1 (2025-09-24)
+
+
+- feat(utility): support to flash bin to encrypted device
+- feat: add runners database
+- feat: add default gen part tool
+- feat(utility): support to check flash enc from bin path
+- fix: get baud from bin path sdkconfig file
+
 ## v0.2.0 (2025-07-16)
 
 

{esp_test_utils-0.2.0 → esp_test_utils-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: esp-test-utils
-Version: 0.2.0
+Version: 0.2.1
 Summary: ESP Test Utils
 Author-email: Chen Yudong <chenyudong@espressif.com>
 License:                              Apache License
@@ -228,7 +228,9 @@ Requires-Dist: pyserial
 Requires-Dist: PyYAML
 Requires-Dist: pexpect
 Requires-Dist: pyusb
+Requires-Dist: esptool
 Requires-Dist: packaging
+Requires-Dist: sqlalchemy
 Requires-Dist: typing_extensions; python_version < "3.11"
 Provides-Extra: idfci
 Requires-Dist: pyecharts; extra == "idfci"

{esp_test_utils-0.2.0 → esp_test_utils-0.2.1}/esp_test_utils.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: esp-test-utils
-Version: 0.2.0
+Version: 0.2.1
 Summary: ESP Test Utils
 Author-email: Chen Yudong <chenyudong@espressif.com>
 License:                              Apache License
@@ -228,7 +228,9 @@ Requires-Dist: pyserial
 Requires-Dist: PyYAML
 Requires-Dist: pexpect
 Requires-Dist: pyusb
+Requires-Dist: esptool
 Requires-Dist: packaging
+Requires-Dist: sqlalchemy
 Requires-Dist: typing_extensions; python_version < "3.11"
 Provides-Extra: idfci
 Requires-Dist: pyecharts; extra == "idfci"

{esp_test_utils-0.2.0 → esp_test_utils-0.2.1}/esp_test_utils.egg-info/SOURCES.txt

@@ -44,6 +44,8 @@ esptest/common/timestamp.py
 esptest/config/__init__.py
 esptest/config/default_config.py
 esptest/config/env_config.py
+esptest/db/__init__.py
+esptest/db/runners.py
 esptest/devices/__init__.py
 esptest/devices/attenuator.py
 esptest/devices/esp_serial.py
@@ -74,7 +76,9 @@ esptest/scripts/monitor.py
 esptest/scripts/set_att.py
 esptest/tools/copy_bin.py
 esptest/tools/download_bin.py
+esptest/tools/http_download.py
 esptest/tools/pip_check.py
+esptest/utility/gen_esp32part.py
 esptest/utility/parse_bin_path.py
 example/jap_test.py
 example/restart_test.py
@@ -86,6 +90,7 @@ tests/test_common.py
 tests/test_import.py
 tests/basic/test_decorators.py
 tests/basic/test_network.py
+tests/db/test_db_runners.py
 tests/esp_console/conftest.py
 tests/esp_console/test_WifiCmd.py
 tests/esp_console/_files/wifi_cmd_connected_1.log
@@ -97,5 +102,9 @@ tests/iperf_utility/_files/dut_iperf_rx1.log
 tests/iperf_utility/_files/dut_iperf_rx2.log
 tests/iperf_utility/_files/pc_iperf_rx.log
 tests/iperf_utility/_files/pc_iperf_rx2.log
+tests/tools/test_download_file.py
 tests/tools/test_pip_check.py
+tests/utility/test_parse_bin_path.py
+tests/utility/_files/test-bin.zip
+tests/utility/_files/test-get-baud/ESP32AT-V4.1.1.0/sdkconfig
 tools/ci/check_dev_version.py

{esp_test_utils-0.2.0 → esp_test_utils-0.2.1}/esp_test_utils.egg-info/requires.txt

@@ -3,7 +3,9 @@ pyserial
 PyYAML
 pexpect
 pyusb
+esptool
 packaging
+sqlalchemy
 
 [:python_version < "3.11"]
 typing_extensions

{esp_test_utils-0.2.0 → esp_test_utils-0.2.1}/esptest/adapter/dut/dut_base.py

@@ -27,7 +27,8 @@ DEFAULT_SERIAL_CONFIGS = {'timeout': 0.005}
 class DutConfig:
     name: str = ''  # default = dut name / port name
     device: str = ''  # log serial device, eg: '/dev/ttyUSB0', 'COM3', etc.
-    baudrate: int = 0  # 0: get from bin path or 115200
+    baudrate: int = 0  # console baudrate, 0: get from bin path or 115200
+    baudrate_from_bin_path: bool = True  # always get baudrate from bin path if bin_path is set
     serial_configs: t.Optional[t.Dict[str, t.Any]] = None  # serial configs, eg: {'bytesize': 8, 'timeout': 0.1}
     # capabilities
     support_esptool: bool = False  # esp port or serial port
@@ -71,10 +72,14 @@ class DutConfig:
         # bin_path and get variables from bin path
         if self.bin_path:
             self.bin_path = Path(self.bin_path).expanduser().resolve()
-            self.esptool_stub = ParseBinPath(self.bin_path).stub
-            self.esptool_chip = ParseBinPath(self.bin_path).chip
-            if not self.baudrate:
-                self.baudrate = get_baud_from_bin_path(self.bin_path) or 115200
+            parsed_bin = ParseBinPath(self.bin_path)
+            self.esptool_stub = parsed_bin.stub
+            self.esptool_chip = parsed_bin.chip
+            if self.baudrate_from_bin_path:
+                # baudrate from bin path is much reliable than the specified one for uart0
+                self.baudrate = get_baud_from_bin_path(self.bin_path) or self.baudrate or 115200
+        if not self.baudrate:
+            self.baudrate = 115200  # set default baudrate to 115200
         # download device
         if not self.download_device:
             self.download_device = self.device

esp_test_utils-0.2.1/esptest/db/runners.py

@@ -0,0 +1,199 @@
+import contextlib
+import csv
+from typing import Any, Dict, Generator, List, Optional
+
+from sqlalchemy import JSON, Engine, ForeignKey, Integer, String
+from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column, relationship
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+class User(Base):
+    __tablename__ = 'user'
+
+    id: Mapped[int] = mapped_column(Integer, autoincrement=True, primary_key=True)
+    name: Mapped[str] = mapped_column(String(50), unique=True)
+    email: Mapped[Optional[str]] = mapped_column(String(50))
+    runners: Mapped[List['Runner']] = relationship(back_populates='user')
+
+
+class Runner(Base):
+    __tablename__ = 'runner'
+
+    id: Mapped[int] = mapped_column(Integer, autoincrement=True, primary_key=True)
+    mac: Mapped[str] = mapped_column(String(50), unique=True)
+    ip: Mapped[str] = mapped_column(String(100), unique=True)
+    vlan: Mapped[Optional[int]] = mapped_column(Integer)
+    switch_info: Mapped[Optional[str]] = mapped_column(String)  # ip-port
+    name: Mapped[str] = mapped_column(String(50))
+    owner: Mapped[Optional[str]] = mapped_column(ForeignKey('user.name'), nullable=True)
+    tags: Mapped[Optional[List[str]]] = mapped_column(JSON)  # only PostgreSQL supports ARRAY(String)
+    raw_data: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON)
+    description: Mapped[Optional[str]] = mapped_column(String)
+    # Link to user table
+    user: Mapped[User] = relationship(back_populates='runners')
+
+    def __str__(self) -> str:
+        tags_str = ','.join(self.tags) if self.tags else ''
+        return (
+            f'Runner('
+            f'id={self.id},'
+            f'mac={self.mac},'
+            f'ip={self.ip},'
+            f'vlan={self.vlan},'
+            f'switch_info={self.switch_info},'
+            f'name={self.name},'
+            f'owner={self.owner},'
+            f'tags="{tags_str}",'
+            f'raw_data={self.raw_data},'
+            f'description="{self.description}"'
+            ')'
+        )
+
+
+class RunnerDB:
+    def __init__(self, engine: Engine) -> None:
+        self.engine = engine
+        self._session: Optional[Session] = None
+        self.__init_db()
+
+    def __init_db(self) -> None:
+        Base.metadata.create_all(self.engine)
+
+    @contextlib.contextmanager
+    def open_session(self, *args: Any, **kwargs: Any) -> Generator['DBSession', None, None]:
+        session = DBSession(self, *args, **kwargs)
+        yield session
+        try:
+            session.commit()
+        except:
+            session.rollback()
+            raise
+        finally:
+            session.close()
+
+    def to_csv(self, csv_file: str) -> None:
+        """
+        Export the runners to a CSV file.
+
+        Args:
+            csv_file: The path to the CSV file.
+            session: The session to use. If None, use the default session.
+        """
+        with self.open_session() as session:
+            runners = session.all_runners()
+            fields = [c.name for c in Runner.__table__.columns]
+            with open(csv_file, 'w', encoding='utf-8') as f:
+                writer = csv.writer(f)
+                writer.writerow(fields)
+                for runner in runners:
+                    writer.writerow([getattr(runner, f) for f in fields])
+
+
+class DBSession(Session):
+    def __init__(self, runner_db: RunnerDB, *args: Any, **kwargs: Any) -> None:
+        self.runner_db = runner_db
+        super().__init__(self.runner_db.engine, *args, **kwargs)
+
+    def all_runners(self) -> List[Runner]:
+        return self.query(Runner).outerjoin(Runner.user).all()
+
+    def all_users(self) -> List[User]:
+        return self.query(User).outerjoin(User.runners).all()
+
+    def get_runners_by_mac(self, mac: str) -> List[Runner]:
+        return self.query(Runner).filter(Runner.mac == mac).all()
+
+    def get_runners_by_ip(self, ip: str) -> List[Runner]:
+        return self.query(Runner).filter(Runner.ip == ip).all()
+
+    def get_users_by_name(self, name: str) -> List[User]:
+        return self.query(User).outerjoin(Runner.user).filter(User.name == name).all()
+
+    def get_users_by_email(self, email: str) -> List[User]:
+        return self.query(User).outerjoin(Runner.user).filter(User.email == email).all()
+
+    def remove_runner(self, runner: Runner) -> None:
+        """
+        Remove a runner.
+
+        Args:
+            runner: The runner to remove.
+            session: The session to use. If None, use the default session.
+        """
+        # session.merge(runner)
+        self.delete(runner)
+        self.flush()
+
+    def add_or_update_runner(self, runner: Runner) -> int:
+        """
+        Add or update a runner.
+
+        Args:
+            runner: The runner to add or update.
+            session: The session to use. If None, use the default session.
+        """
+
+        if runner.id is not None:
+            self.merge(runner)
+        else:
+            if self.get_runners_by_mac(runner.mac) or self.get_runners_by_ip(runner.ip):
+                raise ValueError(f'Runner with mac {runner.mac} or ip {runner.ip} already exists')
+            self.add(runner)
+        self.flush()
+        return runner.id
+
+    def add_or_update_user(self, user: User) -> int:
+        """
+        Add or update an user.
+
+        Args:
+            user: The user to add or update.
+            session: The session to use. If None, use the default session.
+        """
+        if user.id is not None:
+            self.merge(user)
+        else:
+            if self.get_users_by_name(user.name):
+                raise ValueError(f'User with name {user.name} already exists')
+            self.add(user)
+        self.flush()
+        return user.id
+
+
+if __name__ == '__main__':
+    from sqlalchemy import create_engine
+
+    engine1 = create_engine('sqlite:///test.db')
+    # engine = create_engine("mysql+pymysql://user:pass@localhost/testdb")
+    db1 = RunnerDB(engine1)
+    import logging
+
+    logging.basicConfig(level=logging.DEBUG)
+
+    with db1.open_session() as session1:
+        for i in range(5):
+            runner1 = Runner(
+                # id=i+1,
+                mac=f'11:22:33:44:55:6{i}',
+                ip=f'192.168.1.{i}',
+                owner=f'test{i // 2}',
+                name=f'runner00{i}',
+                description=f'test runner 00{i} update 2',
+                tags=['esp32', 'generic', 'eco4'],
+                raw_data={'envs': {'esp32': [f'wifi_iperf_{i}']}, 'os': 'linux'},
+            )
+            session1.add_or_update_runner(runner1)
+        new_user = User(
+            # id=1,
+            name='test0'
+        )
+        session1.add_or_update_user(new_user)
+        session1.commit()
+        users = session1.query(User).outerjoin(User.runners).all()
+        print(users)
+        all_runners = session1.all_runners()
+        print(','.join([str(r) for r in all_runners]))
+    db1.to_csv('runners.csv')

{esp_test_utils-0.2.0 → esp_test_utils-0.2.1}/esptest/scripts/downbin.py

@@ -2,26 +2,34 @@ import argparse
 import logging
 import os
 import re
+import sys
 
 from esptest.devices.serial_tools import get_all_serial_ports
-from esptest.tools.download_bin import download_bin_to_ports
+from esptest.tools.download_bin import bin_path_to_dir, download_bin_to_ports
 
 
 def main() -> None:
-    parser = argparse.ArgumentParser(description='Download bin')
+    usage_string = '%(prog)s [bin_path] [options]'
+    parser = argparse.ArgumentParser(description='Download bin', usage=usage_string)
     parser.add_argument('bin_path', type=str, nargs='?', help='esp bin path, default ./build')
     parser.add_argument('-p', '--ports', type=str, nargs='*', help='download port list')
     parser.add_argument(
         '--range', type=str, help='port list from range (linux), eg: "0-10" equals to "-p ttyUSB0 ttyUSB1 ... ttyUSB10"'
     )
-    parser.add_argument('--all', action='store_true', help='download to all serial ports.')
-    parser.add_argument('--erase-nvs', type=bool, default=True, help='use --erase-nvs=n to skip erase nvs')
-    parser.add_argument('--max-workers', type=int, default=True, help='max download threads')
+    parser.add_argument(
+        '--all', action='store_true', help='download to all serial ports, ignored if "-p/--ports" is specified.'
+    )
+    parser.add_argument('--no-erase-nvs', dest='erase_nvs', action='store_false', help='skip erase nvs')
+    parser.add_argument('--max-workers', type=int, default=0, help='max download threads')
     args = parser.parse_args()
 
     bin_path = args.bin_path or './build'
     if not os.path.isdir(bin_path):
-        raise ValueError(f'Can not find bin_path: {bin_path}')
+        try:
+            bin_path = bin_path_to_dir(bin_path)
+        except Exception as e:  # pylint: disable=broad-except
+            logging.exception(f'Invalid bin path {bin_path} : {str(e)}')
+            sys.exit(1)
 
     ports = []
     if args.ports:
@@ -37,9 +45,13 @@ def main() -> None:
         ports = [os.getenv('ESPPORT') or '/dev/ttyUSB0']
     assert isinstance(ports, list)
 
-    logging.critical(f'Download {bin_path} to {ports}')
-    download_bin_to_ports(bin_path, ports, args.erase_nvs, args.max_workers)
+    try:
+        download_bin_to_ports(bin_path, ports, args.erase_nvs, args.max_workers)
+    except RuntimeError as e:
+        logging.error(str(e))
+        sys.exit(1)
 
 
 if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(message)s')
     main()

{esp_test_utils-0.2.0 → esp_test_utils-0.2.1}/esptest/tools/download_bin.py

@@ -1,13 +1,18 @@
 import asyncio
 import concurrent
 import concurrent.futures
+import os
+import re
 import subprocess
+import tempfile
+import zipfile
 from asyncio.events import AbstractEventLoop
 from functools import lru_cache, partial
 
 import esptest.common.compat_typing as t
 from esptest.devices.serial_tools import compute_serial_port
 from esptest.logger import get_logger
+from esptest.tools.http_download import download_file
 from esptest.utility.parse_bin_path import ParseBinPath
 
 logger = get_logger('download_bin')
@@ -18,6 +23,32 @@ def _get_bin_parser(bin_path: str, parttool: str) -> ParseBinPath:
     return ParseBinPath(bin_path, parttool)
 
 
+@lru_cache()
+def _tmp_dir() -> str:
+    return tempfile.mkdtemp()
+
+
+@lru_cache()
+def bin_path_to_dir(bin_path: str) -> str:
+    bin_hash = hash(bin_path)
+    bin_hash_name = os.path.basename(bin_path)
+    if bin_path.startswith('http'):
+        assert bin_path.endswith('.zip')  # for now only support zip from url
+        new_bin_path = os.path.join(_tmp_dir(), f'{bin_hash}', bin_hash_name)
+        os.makedirs(os.path.dirname(new_bin_path), exist_ok=True)
+        download_file(bin_path, new_bin_path)
+        bin_path = new_bin_path
+    if bin_path.endswith('.zip'):
+        new_bin_path = os.path.join(_tmp_dir(), f'{bin_hash}', bin_hash_name.removesuffix('.zip'))
+        os.makedirs(new_bin_path, exist_ok=True)
+        with zipfile.ZipFile(bin_path, 'r') as zip_ref:
+            zip_ref.extractall(new_bin_path)
+        bin_path = new_bin_path
+    if 'partition_table' not in os.listdir(bin_path):
+        logger.warning('Can not find partition_table from bin_path, maybe invalid!')
+    return bin_path
+
+
 def _filter_esptool_log(log: str) -> str:
     lines = log.splitlines(keepends=True)
     new_log = ''
@@ -36,6 +67,7 @@ def _filter_esptool_log(log: str) -> str:
 class DownBinTool:
     # RETRY_CNT = 2
     DEFAULT_BAUD_LIST = [921600, 460800]
+    FLASH_CRYPT_CNT_PATTERN = re.compile(r'(?:FLASH_CRYPT_CNT|SPI_BOOT_CRYPT_CNT).*\(0b([01]+)')
 
     def __init__(
         self,
@@ -54,11 +86,29 @@ class DownBinTool:
         else:
             self.baud_list = baud
         self.esptool = esptool or 'python -m esptool'
+        self.espefuse = self.esptool.replace('esptool', 'espefuse')
         self.erase_nvs = erase_nvs
         self.bin_parser = _get_bin_parser(bin_path, parttool)
         self.force_no_stub = force_no_stub
 
+    def check_flash_encrypted(self, efuse_summary: str) -> bool:
+        match = self.FLASH_CRYPT_CNT_PATTERN.search(efuse_summary)
+        if match:
+            return match.group(1).count('1') % 2 == 1
+        return False
+
     def download(self) -> None:
+        efuse_cmd = self.espefuse.split()
+        try:
+            summary = subprocess.check_output(
+                efuse_cmd + ['--port', self.port, 'summary'], stderr=subprocess.STDOUT, text=True
+            )
+        except subprocess.CalledProcessError as err:
+            logger.error(err.output)
+            raise RuntimeError(f'Failed to get efuse information from {self.port}') from err
+
+        enc_indicator = ' [encrypted]' if self.check_flash_encrypted(summary) else ''
+
         download_log = ''
         for baud in self.baud_list:
             args = self.esptool.split()
@@ -66,19 +116,18 @@ class DownBinTool:
                 args += ['--no-stub']
             args += ['-p', self.port]
             args += ['-b', f'{baud}']
-            args += self.bin_parser.flash_bin_args(erase_nvs=self.erase_nvs)
+            args += self.bin_parser.flash_bin_args(erase_nvs=self.erase_nvs, encrypted=bool(enc_indicator))
 
-            logger.critical(f'Downloading {self.port}@{baud}: {self.bin_path}')
+            logger.info(f'Downloading {self.port}@{baud}{enc_indicator}: {self.bin_path}')
             # get return code rather than check
             ret = subprocess.run(args, capture_output=True, text=True, check=False)
             if ret.returncode == 0:
                 return  # succeed
             # failed
-            download_log = f'esptool cmd failed ({ret.returncode}): ' + ' '.join(args)
+            download_log += f'esptool cmd failed ({ret.returncode}): ' + ' '.join(args)
             download_log += f'\nDownload failed: [{self.port}@{baud}]\n'
             esptool_msg = ret.stdout + ret.stderr
             download_log += f'esptool output: {_filter_esptool_log(esptool_msg)}'
-            logger.debug(download_log)
         logger.error(download_log)
         raise RuntimeError(f'Failed to download Bin to {self.port}')
 

esp_test_utils-0.2.1/esptest/tools/http_download.py

@@ -0,0 +1,61 @@
+import logging
+import os
+import sys
+import urllib.request
+from typing import Optional
+
+
+def _progress(downloaded: int, total_size: int) -> None:
+    if total_size > 0:
+        percent = min(downloaded / total_size * 100, 100)
+        bar_len = 50
+        filled_len = int(bar_len * downloaded // total_size)
+        progress_bar = '█' * filled_len + '-' * (bar_len - filled_len)
+        sys.stdout.write(f'\r[{progress_bar}] {percent:6.1f}%')
+        sys.stdout.flush()
+    else:
+        # show downloaded size if no total_size
+        sys.stdout.write(f'\rDownloaded {downloaded} bytes')
+        sys.stdout.flush()
+
+
+def download_file(url: str, local_filename: str, timeout: Optional[float] = None, progress: bool = True) -> None:
+    """
+    Download a file from a URL.
+
+    The optional *timeout* parameter specifies a timeout in seconds for
+    blocking operations like the connection attempt (if not specified, the
+    global default timeout setting will be used). This only works for HTTP,
+    HTTPS and FTP connections.
+
+    Args:
+        url: The URL of the file to download.
+        local_filename: The local filename to save the downloaded file.
+        progress: Whether to show the download progress.
+    """
+    if os.path.exists(local_filename):
+        os.remove(local_filename)
+    try:
+        logging.info(f'Downloading {url} -> {local_filename}')
+        with urllib.request.urlopen(url, timeout=timeout) as response, open(local_filename, 'wb') as out_file:
+            total_length = int(response.getheader('Content-Length') or '0')
+            downloaded = 0
+            block_size = 8192
+            while True:
+                chunk = response.read(block_size)
+                if not chunk:
+                    break
+                out_file.write(chunk)
+                downloaded += len(chunk)
+                if progress:
+                    _progress(downloaded, total_length)
+            if progress:
+                sys.stdout.write('\n')
+                sys.stdout.flush()
+            if total_length and total_length != downloaded:
+                logging.error(f'Download {url} failed! maybe url timeout.')
+                raise OSError(f'Download {url} failed: total_length {total_length} != downloaded {downloaded}')
+        logging.info('Download complete!')
+    except OSError as e:
+        logging.error(f'Download {url} failed: {str(e)}')
+        raise e