instrument-data-mcp/scripts/smoke_test.py

from __future__ import annotations

import argparse
import json
from typing import Any

from instrument_config_mcp.config_api import (
    list_device_types,
    list_locations,
    list_meter_types,
    list_system_tree,
    list_systems,
    search_devices,
    search_meters,
)


def _parse_bool(raw: str) -> bool:
    text = str(raw).strip().lower()
    if text in {"1", "true", "yes", "y", "on"}:
        return True
    if text in {"0", "false", "no", "n", "off"}:
        return False
    raise argparse.ArgumentTypeError(f"invalid boolean value: {raw}")


def _print(payload: Any) -> None:
    print(json.dumps(payload, ensure_ascii=False, indent=2))


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Smoke test for instrument config MCP upstream APIs")
    subparsers = parser.add_subparsers(dest="command", required=True)

    p = subparsers.add_parser("list-locations")
    p.add_argument("--project-key", required=True)
    p.add_argument("--keyword", default="")
    p.add_argument("--page-size", type=int, default=100)
    p.add_argument("--page-num", type=int, default=1)

    p = subparsers.add_parser("list-system-tree")
    p.add_argument("--project-key", required=True)

    p = subparsers.add_parser("list-systems")
    p.add_argument("--project-key", required=True)
    p.add_argument("--page-size", type=int, default=100)
    p.add_argument("--page-num", type=int, default=1)
    p.add_argument("--system-type-id", type=int, default=0)
    p.add_argument("--show-below", type=_parse_bool, default=True)

    p = subparsers.add_parser("list-device-types")
    p.add_argument("--project-key", required=True)

    p = subparsers.add_parser("list-meter-types")
    p.add_argument("--project-key", required=True)

    p = subparsers.add_parser("search-devices")
    p.add_argument("--project-key", required=True)
    p.add_argument("--page-size", type=int, default=100)
    p.add_argument("--page-num", type=int, default=1)
    p.add_argument("--keyword", default="")
    p.add_argument("--location-id", type=int, default=0)
    p.add_argument("--show-below", type=_parse_bool, default=True)
    p.add_argument("--system-ids", type=int, nargs="*", default=[])
    p.add_argument("--device-type-ids", type=int, nargs="*", default=[])

    p = subparsers.add_parser("search-meters")
    p.add_argument("--project-key", required=True)
    p.add_argument("--page-size", type=int, default=100)
    p.add_argument("--page-num", type=int, default=1)
    p.add_argument("--keyword", default="")
    p.add_argument("--location-id", type=int, default=0)
    p.add_argument("--show-below", type=_parse_bool, default=True)
    p.add_argument("--meter-type-id", type=int, default=0)
    p.add_argument("--measurement-location-ids", type=int, nargs="*", default=[])
    p.add_argument("--measurement-system-ids", type=int, nargs="*", default=[])
    p.add_argument("--measurement-device-type-ids", type=int, nargs="*", default=[])
    p.add_argument("--status", type=int)

    return parser


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    if args.command == "list-locations":
        _print(list_locations(args.project_key, keyword=args.keyword, page_size=args.page_size, page_num=args.page_num))
        return

    if args.command == "list-system-tree":
        _print(list_system_tree(args.project_key))
        return

    if args.command == "list-systems":
        _print(
            list_systems(
                args.project_key,
                page_size=args.page_size,
                page_num=args.page_num,
                system_type_id=args.system_type_id,
                show_below=args.show_below,
            )
        )
        return

    if args.command == "list-device-types":
        _print(list_device_types(args.project_key))
        return

    if args.command == "list-meter-types":
        _print(list_meter_types(args.project_key))
        return

    if args.command == "search-devices":
        _print(
            search_devices(
                args.project_key,
                page_size=args.page_size,
                page_num=args.page_num,
                keyword=args.keyword,
                location_id=args.location_id,
                show_below=args.show_below,
                system_ids=args.system_ids,
                device_type_ids=args.device_type_ids,
            )
        )
        return

    if args.command == "search-meters":
        _print(
            search_meters(
                args.project_key,
                page_size=args.page_size,
                page_num=args.page_num,
                keyword=args.keyword,
                location_id=args.location_id,
                show_below=args.show_below,
                meter_type_id=args.meter_type_id,
                measurement_location_ids=args.measurement_location_ids,
                measurement_system_ids=args.measurement_system_ids,
                measurement_device_type_ids=args.measurement_device_type_ids,
                status=args.status,
            )
        )
        return

    raise ValueError(f"unsupported command: {args.command}")


if __name__ == "__main__":
    main()