diff --git a/README.md b/README.md index 6a29b0bd..3f192476 100644 --- a/README.md +++ b/README.md @@ -84,7 +84,7 @@ AssetOpsBench is a **unified framework for developing, orchestrating, and evalua | MCP Servers | Important tools | |---|---| -| **IoT** | `sites`, `asset_ids`, `assets` | +| **IoT** | `sites`, `asset_ids`, `asset_detail`, `assets`, `find_assets_by_sensors`, `installed_sensors`, `measured_sensors` | | **FMSR** | `get_failure_modes`, `generate_failure_modes`, `add_failure_modes`, `generate_failure_mode_sensor_mapping` | | **TSFM** | `forecasting`, `timeseries_anomaly_detection` | | **WO** | `get_work_order_distribution`, `predict_next_work_order`, ... | diff --git a/docs/mcp-servers.md b/docs/mcp-servers.md index 989c484b..5fe34c94 100644 --- a/docs/mcp-servers.md +++ b/docs/mcp-servers.md @@ -11,15 +11,18 @@ Six FastMCP servers expose the AssetOpsBench domain logic. Each is a standalone - [tsfm — Time Series Foundation Model](#tsfm--time-series-foundation-model) - [vibration — Vibration Diagnostics](#vibration--vibration-diagnostics) -## iot — IoT Asset Registry +## iot — IoT Asset Registry and Telemetry Records The IoT server reads from the asset **registry** (`ASSET_DBNAME`, default `asset`, loaded from -`asset_profile_sample.json`). It exposes registry discovery tools while the broader IoT tool surface -is being rebuilt: `sites()` for site names, `asset_ids()` for bare `assetnum` values, and `assets()` -for registry metadata with optional `assettype` filtering. +`asset_profile_sample.json`) and IoT telemetry records (`IOT_DBNAME`, default `iot`). It exposes +registry discovery tools while the broader IoT tool surface is being rebuilt: `sites()` for site +names, `asset_ids()` for bare `assetnum` values, `asset_detail()` for one asset's registry +details, `assets()` for registry metadata with optional `assettype` filtering, +`find_assets_by_sensors()` to find assets by installed or measured sensors, `installed_sensors()` +for registry sensor inventory, and `measured_sensors()` for telemetry fields observed in records. **Path:** `src/servers/iot/main.py` -**Requires:** CouchDB (`COUCHDB_URL`, `COUCHDB_USERNAME`, `COUCHDB_PASSWORD`, `ASSET_DBNAME`) +**Requires:** CouchDB (`COUCHDB_URL`, `COUCHDB_USERNAME`, `COUCHDB_PASSWORD`, `ASSET_DBNAME`, `IOT_DBNAME`) **Sample asset profiles shipped in the `asset` database** (loaded by `src/couchdb/init_data.py`): @@ -35,8 +38,19 @@ Source file: `src/couchdb/scenarios_data/shared/iot/asset_profile_sample.json`. | ----------------- | ------------------------------------------ | ------------------------------------------------------------------------------------------------------------ | | `sites` | - | List known site names from the asset registry, with a default fallback | | `asset_ids` | `site_name` | List bare `assetnum` values registered at a site | +| `asset_detail` | `site_name`, `asset_id` | Return one asset's registry details, including `n_installed_sensors` | +| `find_assets_by_sensors` | `site_name`, `sensors`, `match?`, `substring?`, `source?` | Find assets at a site by installed or measured sensor names | +| `installed_sensors` | `site_name`, `asset_id` | List sensor names installed on an asset according to the asset registry | +| `measured_sensors` | `site_name`, `asset_id` | List measured telemetry fields observed for an asset | | `assets` | `site_name`, `assettype?` | List assets with metadata (assettype, description, vintage, installed sensor count), optionally filtered by assettype | +`find_assets_by_sensors()` searches only assets registered at the requested site. `match="all"` +requires every query sensor to match, while `match="any"` requires at least one. Set +`substring=true` to perform case-insensitive substring matching instead of exact sensor-name +matching. `source="installed"` searches the asset registry sensor inventory; `source="measured"` +searches telemetry fields observed in IoT records. Duplicate query sensors are deduplicated in +the response while preserving their first occurrence order. + ## utilities — Utilities **Path:** `src/servers/utilities/main.py` diff --git a/src/servers/iot/main.py b/src/servers/iot/main.py index 6f69ab56..c734b613 100644 --- a/src/servers/iot/main.py +++ b/src/servers/iot/main.py @@ -1,6 +1,6 @@ import logging import os -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, Iterator, List, Optional, Union import couchdb3 from dotenv import load_dotenv @@ -19,8 +19,21 @@ COUCHDB_URL = os.environ.get("COUCHDB_URL") COUCHDB_USERNAME = os.environ.get("COUCHDB_USERNAME") COUCHDB_PASSWORD = os.environ.get("COUCHDB_PASSWORD") +IOT_DBNAME = os.environ.get("IOT_DBNAME", "iot") ASSET_DBNAME = os.environ.get("ASSET_DBNAME", "asset") +try: + iot_db = couchdb3.Database( + IOT_DBNAME, + url=COUCHDB_URL, + user=COUCHDB_USERNAME, + password=COUCHDB_PASSWORD, + ) + logger.info(f"Connected to IoT records database: {IOT_DBNAME}") +except Exception as e: + logger.error(f"Failed to connect to IoT records database: {e}") + iot_db = None + try: asset_db = couchdb3.Database( ASSET_DBNAME, @@ -36,13 +49,18 @@ mcp = FastMCP( "iot", instructions=( - "IoT asset registry tools. Use sites() to discover site names, asset_ids() for bare " - "assetnum values at a site, and assets() for registry metadata with optional " - "assettype filtering." + "IoT asset registry and telemetry record tools. Use sites() to discover site names, " + "asset_ids() for bare assetnum values at a site, asset_detail() for one asset's " + "registry details, assets() for registry metadata with optional assettype filtering, " + "find_assets_by_sensors() to find assets by installed or measured sensors, " + "installed_sensors() for registry sensor inventory, and measured_sensors() for " + "observed telemetry fields." ), ) DEFAULT_SITES = ["MAIN"] +PAGE_SIZE = 1000 +RESERVED_FIELDS = {"_id", "_rev", "asset_id", "timestamp", "dataset", "type", "doctype"} class ErrorResult(BaseModel): @@ -60,6 +78,27 @@ class AssetsResult(BaseModel): message: str +class SensorsResult(BaseModel): + site_name: str + asset_id: str + total_sensors: int + sensors: List[str] + message: str + + +class AssetDetail(BaseModel): + site_name: str + asset_id: str + description: Optional[str] + assettype: Optional[str] + status: Optional[str] + location: Optional[str] + installdate: Optional[str] + vintage: Optional[str] + n_installed_sensors: int + message: str + + class AssetSummary(BaseModel): asset_id: str description: Optional[str] @@ -75,7 +114,74 @@ class AssetsWithMetadataResult(BaseModel): message: str +class AssetSensorMatch(BaseModel): + asset_id: str + matched_sensors: List[str] + + +class FindAssetsResult(BaseModel): + site_name: str + query_sensors: List[str] + match: str + source: str + total_assets: int + assets: List[AssetSensorMatch] + message: str + + _registry_sites_cache: Optional[List[str]] = None +_sensor_list_cache: Dict[str, List[str]] = {} + + +def _iter_records( + selector: Dict[str, Any], + fields: Optional[List[str]] = None, + sort: Optional[List[Dict[str, str]]] = None, + page_size: int = PAGE_SIZE, +) -> Iterator[Dict[str, Any]]: + """Yield telemetry records matching a selector across paged database results.""" + if not iot_db: + return + if sort is None: + sort = [{"asset_id": "asc"}, {"timestamp": "asc"}] + bookmark: Optional[str] = None + while True: + kwargs: Dict[str, Any] = {"limit": page_size, "sort": sort} + if fields is not None: + kwargs["fields"] = fields + if bookmark is not None: + kwargs["bookmark"] = bookmark + res = iot_db.find(selector, **kwargs) + docs = res.get("docs", []) + if not docs: + break + for doc in docs: + yield doc + bookmark = res.get("bookmark") + if bookmark is None or len(docs) < page_size: + break + + +def get_sensor_list(asset_id: str) -> List[str]: + """Return sorted telemetry field names observed across all records for an asset.""" + if asset_id in _sensor_list_cache: + return _sensor_list_cache[asset_id] + if not iot_db: + return [] + try: + found = set() + seen = False + for doc in _iter_records({"asset_id": asset_id}): + seen = True + found.update(key for key in doc.keys() if key not in RESERVED_FIELDS) + if not seen: + return [] + sensors = sorted(found) + _sensor_list_cache[asset_id] = sensors + return sensors + except Exception as e: + logger.error(f"get_sensor_list failed for asset_id {asset_id}: {e}") + return [] def get_registry_sites() -> List[str]: @@ -109,6 +215,38 @@ def _is_known_site(site_name: str) -> bool: return site_name in known_sites() +def _site_asset_ids(site_name: str) -> List[str]: + """Return asset ids registered at a site.""" + if not asset_db: + return [] + try: + res = asset_db.find( + {"siteid": site_name}, + fields=["assetnum"], + limit=100000, + ) + return sorted(doc["assetnum"] for doc in res["docs"]) + except Exception as e: + logger.error(f"_site_asset_ids failed: {e}") + return [] + + +def _installed_sensors(asset_id: str, site_name: Optional[str] = None) -> List[str]: + """Return registry sensor names installed on an asset.""" + if not asset_db: + return [] + try: + selector: Dict[str, Any] = {"assetnum": asset_id} + if site_name is not None: + selector["siteid"] = site_name + res = asset_db.find(selector, fields=["sensors"], limit=1) + docs = res.get("docs", []) + return list(docs[0].get("sensors") or []) if docs else [] + except Exception as e: + logger.error(f"_installed_sensors failed for asset_id {asset_id}: {e}") + return [] + + @mcp.tool(title="List Sites") def sites() -> SitesResult: """List known site names from the asset registry. @@ -156,6 +294,170 @@ def asset_ids(site_name: str) -> Union[AssetsResult, ErrorResult]: return ErrorResult(error=str(e)) +@mcp.tool(title="Get Asset Detail") +def asset_detail(site_name: str, asset_id: str) -> Union[AssetDetail, ErrorResult]: + """Return registry details for one asset. + + This tool reads the asset registry from `asset_db` (`ASSET_DBNAME`, default + `asset`) and returns nameplate fields for one asset record. Use + `installed_sensors()` when you only need the registry sensor inventory. + + Args: + site_name: Exact site id to query, such as `MAIN`. Use `sites()` to + discover valid site ids. + asset_id: Exact asset id from `asset_ids()`, such as `Chiller 6`. + + Returns: + AssetDetail: Contains `site_name`, `asset_id`, `description`, + `assettype`, `status`, `location`, `installdate`, `vintage`, + `n_installed_sensors`, and `message`. + """ + if not _is_known_site(site_name): + return ErrorResult(error=f"unknown site {site_name}") + if not asset_db: + return ErrorResult(error="asset registry database not connected") + + try: + res = asset_db.find( + {"siteid": site_name, "assetnum": asset_id}, + fields=[ + "assetnum", + "description", + "assettype", + "status", + "location", + "installdate", + "vintage", + "sensors", + ], + limit=1, + ) + docs = res.get("docs", []) + if not docs: + return ErrorResult(error=f"unknown asset_id {asset_id} at site {site_name}") + + doc = docs[0] + sensors = list(doc.get("sensors") or []) + n_installed_sensors = len(sensors) + assettype = doc.get("assettype") + vintage = doc.get("vintage") + location = doc.get("location") + + parts = [f"asset {asset_id} is a {assettype or 'asset'}"] + if vintage: + parts.append(f" ({vintage} vintage)") + if location: + parts.append(f" at {location}") + parts.append(f" with {n_installed_sensors} installed sensors.") + + return AssetDetail( + site_name=site_name, + asset_id=doc.get("assetnum", asset_id), + description=doc.get("description"), + assettype=assettype, + status=doc.get("status"), + location=location, + installdate=doc.get("installdate"), + vintage=vintage, + n_installed_sensors=n_installed_sensors, + message="".join(parts), + ) + except Exception as e: + logger.error(f"asset_detail failed: {e}") + return ErrorResult(error=str(e)) + + +@mcp.tool(title="List Measured Sensors") +def measured_sensors( + site_name: str, asset_id: str +) -> Union[SensorsResult, ErrorResult]: + """List telemetry fields actually measured for one asset. + + This tool reads IoT records from `iot_db` (`IOT_DBNAME`, default `iot`). + It scans records matching `asset_id` and returns the union of observed + measurement fields. Use `installed_sensors()` instead when you need the + asset registry inventory. + + Args: + site_name: Exact site id to query, such as `MAIN`. Use `sites()` to + discover valid site ids. + asset_id: Exact asset id from `asset_ids()`, such as `Chiller 6`. + + Returns: + SensorsResult: Contains `site_name`, `asset_id`, `total_sensors`, + `sensors`, and `message`. The `sensors` field is a sorted list of + observed telemetry fields, excluding record metadata such as `_id`, + `_rev`, `asset_id`, `timestamp`, `dataset`, `type`, and `doctype`. + """ + if not _is_known_site(site_name): + return ErrorResult(error=f"unknown site {site_name}") + if not iot_db: + return ErrorResult(error="IoT records database not connected") + + sensor_list = get_sensor_list(asset_id) + if not sensor_list: + return ErrorResult(error=f"unknown asset_id {asset_id} or no sensors found") + + return SensorsResult( + site_name=site_name, + asset_id=asset_id, + total_sensors=len(sensor_list), + sensors=sensor_list, + message=( + f"found {len(sensor_list)} sensors for asset_id {asset_id} " + f"and site_name {site_name}: {', '.join(sensor_list)}." + ), + ) + + +@mcp.tool(title="List Installed Sensors") +def installed_sensors( + site_name: str, asset_id: str +) -> Union[SensorsResult, ErrorResult]: + """List sensor names installed on one asset according to the registry. + + This tool reads the asset registry from `asset_db` (`ASSET_DBNAME`, default + `asset`) and returns the asset record's `sensors` inventory. Use + `measured_sensors()` instead when you need fields actually observed in IoT + telemetry records. + + Args: + site_name: Exact site id to query, such as `MAIN`. Use `sites()` to + discover valid site ids. + asset_id: Exact asset id from `asset_ids()`, such as `Chiller 6`. + + Returns: + SensorsResult: Contains `site_name`, `asset_id`, `total_sensors`, + `sensors`, and `message`. The `sensors` field preserves the registry + sensor inventory order from the asset record. + """ + if not _is_known_site(site_name): + return ErrorResult(error=f"unknown site {site_name}") + if not asset_db: + return ErrorResult(error="asset registry database not connected") + + try: + res = asset_db.find( + {"siteid": site_name, "assetnum": asset_id}, + fields=["assetnum", "sensors"], + limit=1, + ) + docs = res.get("docs", []) + if not docs: + return ErrorResult(error=f"unknown asset_id {asset_id} at site {site_name}") + names = list(docs[0].get("sensors") or []) + return SensorsResult( + site_name=site_name, + asset_id=asset_id, + total_sensors=len(names), + sensors=names, + message=f"{len(names)} sensors installed on {asset_id}: {', '.join(names)}.", + ) + except Exception as e: + logger.error(f"installed_sensors failed: {e}") + return ErrorResult(error=str(e)) + + @mcp.tool(title="List Assets") def assets( site_name: str, assettype: Optional[str] = None @@ -194,7 +496,7 @@ def assets( assettype=doc.get("assettype"), description=doc.get("description"), vintage=doc.get("vintage"), - n_sensors=len(doc.get("sensors", [])), + n_sensors=len(doc.get("sensors") or []), ) for doc in res["docs"] ), @@ -213,6 +515,102 @@ def assets( return ErrorResult(error=str(e)) +@mcp.tool(title="Find Assets By Sensors") +def find_assets_by_sensors( + site_name: str, + sensors: List[str], + match: str = "all", + substring: bool = False, + source: str = "measured", +) -> Union[FindAssetsResult, ErrorResult]: + """Return assets at a site that match the requested sensor names. + + The search is limited to assets registered at `site_name`. Duplicate query + sensors are ignored while preserving the first occurrence order. + + Args: + site_name: Exact site id to query, such as `MAIN`. Use `sites()` to + discover valid site ids. + sensors: One or more sensor names or, when `substring` is true, sensor + name fragments to search for. + match: `all` requires every listed sensor; `any` requires at least one. + substring: If true, match sensor names case-insensitively by substring. + source: `measured` checks telemetry fields; `installed` checks the + asset registry inventory. + + Returns: + FindAssetsResult: Contains the deduplicated query sensors, matching + asset ids, and the concrete sensor names that matched each asset. + """ + if not _is_known_site(site_name): + return ErrorResult(error=f"unknown site {site_name}") + if match not in ("all", "any"): + return ErrorResult(error="match must be 'all' or 'any'") + if source not in ("measured", "installed"): + return ErrorResult(error="source must be 'measured' or 'installed'") + if not sensors: + return ErrorResult(error="provide at least one sensor name") + if not asset_db: + return ErrorResult(error="asset registry database not connected") + if source == "measured" and not iot_db: + return ErrorResult(error="IoT records database not connected") + + query_sensors = list(dict.fromkeys(sensors)) + matches: List[AssetSensorMatch] = [] + for asset_id in _site_asset_ids(site_name): + available = ( + get_sensor_list(asset_id) + if source == "measured" + else _installed_sensors(asset_id, site_name) + ) + + def _hits(sensor_name: str) -> List[str]: + if substring: + sensor_name_lower = sensor_name.lower() + return [ + available_sensor + for available_sensor in available + if sensor_name_lower in available_sensor.lower() + ] + return [ + available_sensor + for available_sensor in available + if available_sensor == sensor_name + ] + + per_query = {sensor_name: _hits(sensor_name) for sensor_name in query_sensors} + satisfied = [ + sensor_name for sensor_name, hits in per_query.items() if hits + ] + ok = ( + len(satisfied) == len(query_sensors) + if match == "all" + else len(satisfied) > 0 + ) + if ok: + matched = sorted( + { + matched_sensor + for hits in per_query.values() + for matched_sensor in hits + } + ) + matches.append( + AssetSensorMatch(asset_id=asset_id, matched_sensors=matched) + ) + + return FindAssetsResult( + site_name=site_name, + query_sensors=query_sensors, + match=match, + source=source, + total_assets=len(matches), + assets=matches, + message=f"{len(matches)} asset(s) at {site_name} match {query_sensors} " + f"(match={match}, substring={substring}, source={source}).", + ) + + def main(): mcp.run(transport="stdio") diff --git a/src/servers/iot/tests/conftest.py b/src/servers/iot/tests/conftest.py index 975d9222..1c758e86 100644 --- a/src/servers/iot/tests/conftest.py +++ b/src/servers/iot/tests/conftest.py @@ -31,6 +31,27 @@ def _couchdb_reachable() -> bool: return False +def _iot_data_reachable() -> bool: + url = os.environ.get("COUCHDB_URL") + if not url: + return False + try: + import couchdb3 + + username = os.environ.get("COUCHDB_USERNAME") + password = os.environ.get("COUCHDB_PASSWORD") + iot_db = couchdb3.Database( + os.environ.get("IOT_DBNAME", "iot"), + url=url, + user=username, + password=password, + ) + records = iot_db.find({"asset_id": "Chiller 6"}, limit=1)["docs"] + return bool(records) + except Exception: + return False + + requires_couchdb = pytest.mark.skipif( not _couchdb_reachable(), reason=( @@ -40,6 +61,15 @@ def _couchdb_reachable() -> bool: ) +requires_iot_db = pytest.mark.skipif( + not _iot_data_reachable(), + reason=( + "IoT sample records database not available " + "(set COUCHDB_URL and load the Chiller 6 telemetry records)" + ), +) + + # --- Fixtures --- @@ -50,21 +80,37 @@ def mock_asset_db(): yield mock +@pytest.fixture +def mock_iot_db(): + """Patch the module-level telemetry records `iot_db` object in main with a mock.""" + with patch("servers.iot.main.iot_db") as mock: + yield mock + + @pytest.fixture def no_asset_db(): - """Patch the module-level `asset_db` to None (simulate disconnected CouchDB).""" + """Patch the module-level `asset_db` to None (simulate disconnected database).""" with patch("servers.iot.main.asset_db", None): yield +@pytest.fixture +def no_iot_db(): + """Patch the module-level telemetry records `iot_db` to None.""" + with patch("servers.iot.main.iot_db", None): + yield + + @pytest.fixture(autouse=True) def clear_iot_caches(): """Clear module-level caches so mocked DB responses do not leak across tests.""" import servers.iot.main as iot_main iot_main._registry_sites_cache = None + iot_main._sensor_list_cache = {} yield iot_main._registry_sites_cache = None + iot_main._sensor_list_cache = {} async def call_tool(mcp_instance, tool_name: str, args: dict) -> dict: diff --git a/src/servers/iot/tests/test_tools.py b/src/servers/iot/tests/test_tools.py index 2b308adc..ff8d9c30 100644 --- a/src/servers/iot/tests/test_tools.py +++ b/src/servers/iot/tests/test_tools.py @@ -3,14 +3,22 @@ import pytest from servers.iot.main import mcp -from .conftest import call_tool, requires_couchdb +from .conftest import call_tool, requires_couchdb, requires_iot_db class TestToolRegistration: @pytest.mark.anyio async def test_registry_tools_are_registered(self): tools = await mcp.list_tools() - assert sorted(tool.name for tool in tools) == ["asset_ids", "assets", "sites"] + assert sorted(tool.name for tool in tools) == [ + "asset_detail", + "asset_ids", + "assets", + "find_assets_by_sensors", + "installed_sensors", + "measured_sensors", + "sites", + ] class TestSites: @@ -67,6 +75,404 @@ async def test_discovery_integration(self): assert data["total_assets"] > 0 +class TestAssetDetail: + @pytest.mark.anyio + async def test_invalid_site(self): + data = await call_tool( + mcp, "asset_detail", {"site_name": "INVALID", "asset_id": "Pump-1"} + ) + assert "error" in data + assert "unknown site" in data["error"] + + @pytest.mark.anyio + async def test_with_mock_asset_db(self, mock_asset_db): + mock_asset_db.find.side_effect = [ + {"docs": [{"siteid": "MAIN"}]}, + { + "docs": [ + { + "assetnum": "Pump-1", + "description": "Main pump", + "assettype": "PUMP", + "status": "OPERATING", + "location": "PUMP-HOUSE", + "installdate": "2024-01-01", + "vintage": "new", + "sensors": ["Pressure", "Temperature"], + } + ] + }, + ] + + data = await call_tool( + mcp, "asset_detail", {"site_name": "MAIN", "asset_id": "Pump-1"} + ) + + assert data == { + "site_name": "MAIN", + "asset_id": "Pump-1", + "description": "Main pump", + "assettype": "PUMP", + "status": "OPERATING", + "location": "PUMP-HOUSE", + "installdate": "2024-01-01", + "vintage": "new", + "n_installed_sensors": 2, + "message": "asset Pump-1 is a PUMP (new vintage) at PUMP-HOUSE with 2 installed sensors.", + } + + @pytest.mark.anyio + async def test_reads_asset_registry_not_iot_db(self, mock_asset_db, mock_iot_db): + mock_asset_db.find.side_effect = [ + {"docs": [{"siteid": "MAIN"}]}, + { + "docs": [ + { + "assetnum": "Pump-1", + "description": "Registry pump", + "assettype": "PUMP", + "status": "OPERATING", + "location": None, + "installdate": None, + "vintage": None, + "sensors": [], + } + ] + }, + ] + mock_iot_db.find.return_value = { + "docs": [ + { + "asset_id": "Pump-1", + "timestamp": "2024-01-01T00:00:00", + "Telemetry Sensor": 42, + } + ] + } + + data = await call_tool( + mcp, "asset_detail", {"site_name": "MAIN", "asset_id": "Pump-1"} + ) + + assert data["description"] == "Registry pump" + assert data["n_installed_sensors"] == 0 + mock_iot_db.find.assert_not_called() + + @pytest.mark.anyio + async def test_db_disconnected(self, no_asset_db): + data = await call_tool( + mcp, "asset_detail", {"site_name": "MAIN", "asset_id": "Pump-1"} + ) + assert "error" in data + assert "not connected" in data["error"].lower() + + @requires_couchdb + @pytest.mark.anyio + async def test_discovery_integration(self): + data = await call_tool( + mcp, "asset_detail", {"site_name": "MAIN", "asset_id": "Chiller 6"} + ) + assert data["asset_id"] == "Chiller 6" + assert data["assettype"] == "CHILLER" + assert data["status"] == "OPERATING" + assert data["n_installed_sensors"] > 0 + + +class TestMeasuredSensors: + @pytest.mark.anyio + async def test_invalid_site(self): + data = await call_tool( + mcp, "measured_sensors", {"site_name": "INVALID", "asset_id": "Pump-1"} + ) + assert "error" in data + assert "unknown site" in data["error"] + + @pytest.mark.anyio + async def test_with_mock_iot_db(self, mock_asset_db, mock_iot_db): + mock_asset_db.find.return_value = {"docs": [{"siteid": "MAIN"}]} + mock_iot_db.find.return_value = { + "docs": [ + { + "_id": "iot:Pump-1:1", + "_rev": "1-abc", + "asset_id": "Pump-1", + "timestamp": "2024-01-01T00:00:00", + "dataset": "iot", + "Pressure": 10, + }, + { + "asset_id": "Pump-1", + "timestamp": "2024-01-01T00:01:00", + "dataset": "iot", + "Temperature": 30, + "Pressure": 11, + }, + ] + } + + data = await call_tool( + mcp, "measured_sensors", {"site_name": "MAIN", "asset_id": "Pump-1"} + ) + + assert data["site_name"] == "MAIN" + assert data["asset_id"] == "Pump-1" + assert data["total_sensors"] == 2 + assert data["sensors"] == ["Pressure", "Temperature"] + + @pytest.mark.anyio + async def test_reads_iot_db_not_asset_registry(self, mock_asset_db, mock_iot_db): + mock_asset_db.find.return_value = { + "docs": [ + { + "siteid": "MAIN", + "assetnum": "Pump-1", + "sensors": ["Registry Sensor"], + } + ] + } + mock_iot_db.find.return_value = { + "docs": [ + { + "asset_id": "Pump-1", + "timestamp": "2024-01-01T00:00:00", + "Telemetry Sensor": 42, + } + ] + } + + data = await call_tool( + mcp, "measured_sensors", {"site_name": "MAIN", "asset_id": "Pump-1"} + ) + + assert data["sensors"] == ["Telemetry Sensor"] + assert mock_asset_db.find.call_count == 1 + mock_asset_db.find.assert_called_once_with( + {"siteid": {"$exists": True}}, + fields=["siteid"], + limit=100000, + ) + mock_iot_db.find.assert_called_once() + + @pytest.mark.anyio + async def test_db_disconnected(self, mock_asset_db, no_iot_db): + mock_asset_db.find.return_value = {"docs": [{"siteid": "MAIN"}]} + data = await call_tool( + mcp, "measured_sensors", {"site_name": "MAIN", "asset_id": "Pump-1"} + ) + assert "error" in data + assert "not connected" in data["error"].lower() + + @requires_iot_db + @pytest.mark.anyio + async def test_discovery_integration(self): + data = await call_tool( + mcp, "measured_sensors", {"site_name": "MAIN", "asset_id": "Chiller 6"} + ) + assert "sensors" in data + assert "Chiller 6 Supply Temperature" in data["sensors"] + assert data["total_sensors"] > 0 + + +class TestInstalledSensors: + @pytest.mark.anyio + async def test_invalid_site(self): + data = await call_tool( + mcp, "installed_sensors", {"site_name": "INVALID", "asset_id": "Pump-1"} + ) + assert "error" in data + assert "unknown site" in data["error"] + + @pytest.mark.anyio + async def test_with_mock_asset_db(self, mock_asset_db): + mock_asset_db.find.side_effect = [ + {"docs": [{"siteid": "MAIN"}]}, + { + "docs": [ + { + "assetnum": "Pump-1", + "sensors": ["Registry Pressure", "Registry Temperature"], + } + ] + }, + ] + + data = await call_tool( + mcp, "installed_sensors", {"site_name": "MAIN", "asset_id": "Pump-1"} + ) + + assert data["site_name"] == "MAIN" + assert data["asset_id"] == "Pump-1" + assert data["total_sensors"] == 2 + assert data["sensors"] == ["Registry Pressure", "Registry Temperature"] + mock_asset_db.find.assert_called_with( + {"siteid": "MAIN", "assetnum": "Pump-1"}, + fields=["assetnum", "sensors"], + limit=1, + ) + + @pytest.mark.anyio + async def test_reads_asset_registry_not_iot_db(self, mock_asset_db, mock_iot_db): + mock_asset_db.find.side_effect = [ + {"docs": [{"siteid": "MAIN"}]}, + { + "docs": [ + { + "assetnum": "Pump-1", + "sensors": ["Registry Sensor"], + } + ] + }, + ] + mock_iot_db.find.return_value = { + "docs": [ + { + "asset_id": "Pump-1", + "timestamp": "2024-01-01T00:00:00", + "Telemetry Sensor": 42, + } + ] + } + + data = await call_tool( + mcp, "installed_sensors", {"site_name": "MAIN", "asset_id": "Pump-1"} + ) + + assert data["sensors"] == ["Registry Sensor"] + mock_iot_db.find.assert_not_called() + + @pytest.mark.anyio + async def test_db_disconnected(self, no_asset_db): + data = await call_tool( + mcp, "installed_sensors", {"site_name": "MAIN", "asset_id": "Pump-1"} + ) + assert "error" in data + assert "not connected" in data["error"].lower() + + @requires_couchdb + @pytest.mark.anyio + async def test_discovery_integration(self): + data = await call_tool( + mcp, "installed_sensors", {"site_name": "MAIN", "asset_id": "Chiller 6"} + ) + assert "sensors" in data + assert "Chiller 6 Oil Pressure" in data["sensors"] + assert data["total_sensors"] > 0 + + +class TestFindAssetsBySensors: + @pytest.mark.anyio + async def test_invalid_site(self): + data = await call_tool( + mcp, + "find_assets_by_sensors", + {"site_name": "INVALID", "sensors": ["Pressure"]}, + ) + assert "error" in data + assert "unknown site" in data["error"] + + @pytest.mark.anyio + async def test_installed_source_exact_match(self, mock_asset_db): + mock_asset_db.find.side_effect = [ + {"docs": [{"siteid": "MAIN"}]}, + {"docs": [{"assetnum": "Fan-2"}, {"assetnum": "Pump-1"}]}, + {"docs": [{"sensors": ["Temperature"]}]}, + {"docs": [{"sensors": ["Pressure", "Temperature"]}]}, + ] + + data = await call_tool( + mcp, + "find_assets_by_sensors", + { + "site_name": "MAIN", + "sensors": ["Pressure"], + "source": "installed", + }, + ) + + assert data["site_name"] == "MAIN" + assert data["query_sensors"] == ["Pressure"] + assert data["match"] == "all" + assert data["source"] == "installed" + assert data["total_assets"] == 1 + assert data["assets"] == [ + {"asset_id": "Pump-1", "matched_sensors": ["Pressure"]} + ] + + @pytest.mark.anyio + async def test_deduplicates_query_sensors_for_all_match(self, mock_asset_db): + mock_asset_db.find.side_effect = [ + {"docs": [{"siteid": "MAIN"}]}, + {"docs": [{"assetnum": "Pump-1"}]}, + {"docs": [{"sensors": ["Pressure", "Temperature"]}]}, + ] + + data = await call_tool( + mcp, + "find_assets_by_sensors", + { + "site_name": "MAIN", + "sensors": ["Pressure", "Pressure"], + "source": "installed", + }, + ) + + assert data["query_sensors"] == ["Pressure"] + assert data["total_assets"] == 1 + assert data["assets"] == [ + {"asset_id": "Pump-1", "matched_sensors": ["Pressure"]} + ] + + @pytest.mark.anyio + async def test_measured_source_substring_match(self, mock_asset_db, mock_iot_db): + mock_asset_db.find.side_effect = [ + {"docs": [{"siteid": "MAIN"}]}, + {"docs": [{"assetnum": "Compressor-1"}, {"assetnum": "Pump-1"}]}, + ] + + def find_records(selector, **kwargs): + asset_id = selector["asset_id"] + if asset_id == "Compressor-1": + return { + "docs": [ + { + "asset_id": "Compressor-1", + "timestamp": "2024-01-01T00:00:00", + "Oil Pressure": 12, + } + ] + } + if asset_id == "Pump-1": + return { + "docs": [ + { + "asset_id": "Pump-1", + "timestamp": "2024-01-01T00:00:00", + "Discharge Pressure": 42, + "Flow": 4, + } + ] + } + return {"docs": []} + + mock_iot_db.find.side_effect = find_records + + data = await call_tool( + mcp, + "find_assets_by_sensors", + { + "site_name": "MAIN", + "sensors": ["pressure"], + "substring": True, + }, + ) + + assert data["total_assets"] == 2 + assert data["assets"] == [ + {"asset_id": "Compressor-1", "matched_sensors": ["Oil Pressure"]}, + {"asset_id": "Pump-1", "matched_sensors": ["Discharge Pressure"]}, + ] + + class TestAssets: @pytest.mark.anyio async def test_invalid_site(self):