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 (
    get_topology,
    get_topology_data,
    list_device_types,
    list_locations,
    list_meter_types,
    list_system_tree,
    list_systems,
    search_devices,
    search_meters,
    search_points,
)
from instrument_config_mcp.topology_cache import (
    find_topology_context,
    get_topology_group_config,
    get_topology_node,
    list_topologies,
    list_topology_groups,
)


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)

    p = subparsers.add_parser("search-points")
    p.add_argument("--project-key", required=True)
    p.add_argument("--id", type=int, required=True)
    p.add_argument("--page-size", type=int, default=100)
    p.add_argument("--page-num", type=int, default=1)

    p = subparsers.add_parser("get-topology")
    p.add_argument("--project-key", required=True)
    p.add_argument("--id", type=int, required=True)

    p = subparsers.add_parser("get-topology-data")
    p.add_argument("--project-key", required=True)
    p.add_argument("--id", type=int, required=True)
    p.add_argument("--display", choices=["instant", "accu"], required=True)
    p.add_argument("--accu-step", type=int)
    p.add_argument("--ts", type=int)

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

    p = subparsers.add_parser("topology-list")
    p.add_argument("--project-key", required=True)
    p.add_argument("--group-id", type=int)
    p.add_argument("--object-type-code", type=int)

    p = subparsers.add_parser("topology-get-node")
    p.add_argument("--project-key", required=True)
    p.add_argument("--topology-id", type=int, required=True)
    p.add_argument("--node-id", default="root")
    p.add_argument("--include-siblings", type=_parse_bool, default=True)
    p.add_argument("--include-children", type=_parse_bool, default=True)

    p = subparsers.add_parser("topology-get-group-config")
    p.add_argument("--project-key", required=True)
    p.add_argument("--topology-id", type=int, required=True)

    p = subparsers.add_parser("topology-find-context")
    p.add_argument("--project-key", required=True)
    p.add_argument("--entity-type", choices=["meter", "device"], required=True)
    p.add_argument("--entity-id", type=int, required=True)
    p.add_argument("--topology-id", type=int)
    p.add_argument("--include-siblings", type=_parse_bool, default=True)
    p.add_argument("--ancestor-depth", type=int, default=5)
    p.add_argument("--descendant-depth", type=int, default=2)

    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

    if args.command == "search-points":
        _print(
            search_points(
                args.project_key,
                id=args.id,
                page_size=args.page_size,
                page_num=args.page_num,
            )
        )
        return

    if args.command == "get-topology":
        _print(get_topology(args.project_key, id=args.id))
        return

    if args.command == "get-topology-data":
        _print(
            get_topology_data(
                args.project_key,
                id=args.id,
                display=args.display,
                accu_step=args.accu_step,
                ts=args.ts,
            )
        )
        return

    if args.command == "topology-group-list":
        _print(list_topology_groups(args.project_key))
        return

    if args.command == "topology-list":
        _print(
            list_topologies(
                args.project_key,
                group_id=args.group_id,
                object_type_code=args.object_type_code,
            )
        )
        return

    if args.command == "topology-get-node":
        _print(
            get_topology_node(
                args.project_key,
                args.topology_id,
                args.node_id,
                include_siblings=args.include_siblings,
                include_children=args.include_children,
            )
        )
        return

    if args.command == "topology-get-group-config":
        _print(get_topology_group_config(args.project_key, args.topology_id))
        return

    if args.command == "topology-find-context":
        _print(
            find_topology_context(
                args.project_key,
                args.entity_type,
                args.entity_id,
                topology_id=args.topology_id,
                include_siblings=args.include_siblings,
                ancestor_depth=args.ancestor_depth,
                descendant_depth=args.descendant_depth,
            )
        )
        return

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


if __name__ == "__main__":
    main()