vent_agent_2.0/tests/test_main_fan_realtime.py

"""三个新入口的请求契约与原始数据归档；返回字段使用合成样例，不代表真实平台schema。"""
from contextlib import asynccontextmanager
from copy import deepcopy
import json
from types import SimpleNamespace

import httpx
import pytest
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
from langchain_core.messages import AIMessageChunk
from langchain_core.outputs import ChatGenerationChunk

import config
from agents import agent_factory, agent_registry
from agents.workspace import session_root
from db import chat_store, main_fan_store
from extensions.main_fan import exploration, gateway_queries
from extensions.main_fan.exploration_tools import (query_fan_sensor_realtime,
    query_fan_equipment_realtime, query_fan_equipment_timeseries)
from extensions.shaanxi.client import GatewayClient
from extensions.shaanxi.settings import GatewaySettings
from runtime.environment import RuntimeEnvironment
from runtime.execution import execute
from runtime.python_sandbox import capability, digest, workspace_file
from sandbox_agent_helpers import tool_call

TYPES = ['0005', '0013', '1001', '1010', '0001,0011,0043', '0002']
NODE = '6108310185903000MN30102100000000000064'
CURVE_ARGS = {'dept_id': '1237', 'node_id': NODE, 'start_time': '2026-10-20 06:20:49',
              'end_time': '2026-10-20 09:20:49'}
SENSORS = {'code': 200, 'total': 1, 'rows': [{'nodeId': 'synthetic-switch', 'currentData': '开',
            'sensorTypeId': '1010', 'nodePlace': '合成主通1一级'}]}
EQUIPMENT = {'code': 200, 'total': '3', 'rows': [
    {'nodeId': 'synthetic-flow', 'label': '合成风量', 'currentValue': '3800', 'unit': 'm³/min',
     'device': {'id': 'synthetic-fan', 'name': '合成测试风机'}},
    {'nodeId': 'synthetic-bearing', 'label': '合成电机轴承温度', 'currentValue': '46.2', 'unit': '℃'}]}
CURVE = {'code': 200, 'data': {'metadata': {'nodeId': NODE, 'unit': '℃'},
         'axis': ['2026-10-20 06:20:49', '2026-10-20 09:20:49'], 'values': [45.2, 46.2]}}


@pytest.fixture
def workbench(monkeypatch, tmp_path):
    for directory in ('skills', 'db'):
        (tmp_path / directory).symlink_to(config.ROOT / directory, target_is_directory=True)
    monkeypatch.setattr(config, 'ROOT', tmp_path)
    for key, value in {'MAIN_FAN_TARGETS_FILE': '', 'MAIN_FAN_RULES_FILE': '',
                       'SHANXI_GATEWAY_ENABLED': '1', 'MAIN_FAN_AGENT_CODE': 'main_fan_equivalence_agent',
                       'MAIN_FAN_REALTIME_MAX_PAGE_SIZE': '1000', 'MAIN_FAN_SERIES_MAX_HOURS': '24'}.items():
        monkeypatch.setenv(key, value)
    user = 'realtime-test'
    session = chat_store.create_session(user)['id']
    request = {'mine_code': 'synthetic-mine', 'mine_name': '合成接口测试矿', 'use_gateway': True}
    prepared = exploration.prepare(request, user, session)
    runtime = SimpleNamespace(config={'configurable': {'workspace_user': user, 'workspace_session': session}})
    return SimpleNamespace(user=user, session=session, request=request, identifier=prepared['workflow_id'],
                           runtime=runtime, workspace=session_root(user, session) / 'workspace')


@asynccontextmanager
async def gateway(monkeypatch, response=SENSORS, *, proxy_status=200):
    calls = []
    def handler(request):
        assert request.url.path == '/internal/monitor/proxy'
        assert request.headers['X-Agent-Code'] == 'main_fan_equivalence_agent'
        assert request.headers['Authorization'] == 'Bearer realtime-test-secret'
        payload = json.loads(request.content)
        calls.append(payload)
        body = response(payload) if callable(response) else response
        return httpx.Response(200, json={'code': 200, 'success': True,
            'data': {'status_code': proxy_status, 'body': deepcopy(body)}})
    client = GatewayClient(GatewaySettings('https://gateway.invalid', 'realtime-test-secret'),
                           transport=httpx.MockTransport(handler))
    monkeypatch.setattr(gateway_queries, 'get_service', lambda name: client)
    try:
        yield calls
    finally:
        await client.close()


async def sensor_call(workbench, **overrides):
    args = {'workflow_id': workbench.identifier, 'dept_id': '1117', 'runtime': workbench.runtime, **overrides}
    return await query_fan_sensor_realtime(**args)


async def equipment_call(workbench, **overrides):
    args = {'workflow_id': workbench.identifier, 'dept_id': '1219', 'runtime': workbench.runtime, **overrides}
    return await query_fan_equipment_realtime(**args)


def read_data(workbench, result):
    return json.loads(workspace_file(workbench.user, workbench.session, result['data_path']).read_text())


@pytest.mark.asyncio
async def test_sensor_indexed_filter_matches_curl_and_preserves_complete_payload(workbench, monkeypatch):
    async with gateway(monkeypatch) as calls:
        result = await sensor_call(workbench, sensor_type_list=TYPES)
    assert result['ok'], result
    assert calls[0]['path'] == '/aqjk/node/getNodeRealtimeDetails'
    assert calls[0]['params'] == [['pageSize', 100], ['pageNum', 1], ['deptId', '1117'],
        *[[f'sensorTypeList[{index}]', value] for index, value in enumerate(TYPES)]]
    assert result['row_count'] == 1 and result['total'] == 1 and result['has_more'] is False
    assert read_data(workbench, result)['payload'] == SENSORS
    assert result['rows_pointer'] == '/payload/rows'
    assert result['kind'] == 'sensor_realtime'
    assert 'realtime-test-secret' not in json.dumps(read_data(workbench, result))


@pytest.mark.asyncio
async def test_unfiltered_sensor_query_and_explicit_last_page(workbench, monkeypatch):
    body = {'code': 200, 'data': {'rows': [{'id': 'three'}, {'id': 'four'}], 'total': '4'}}
    async with gateway(monkeypatch, body) as calls:
        result = await sensor_call(workbench, page_num=2, page_size=2)
    assert calls[0]['params'] == [['pageSize', 2], ['pageNum', 2], ['deptId', '1117']]
    assert result['row_count'] == 2 and result['has_more'] is False and result['next_page'] is None
    assert result['rows_pointer'] == '/payload/data/rows'


@pytest.mark.asyncio
async def test_equipment_system_pagination_and_original_units(workbench, monkeypatch):
    def response(payload):
        body = deepcopy(EQUIPMENT)
        if payload['params']['pageNum'] == '2':
            body['rows'] = [{'nodeId': 'synthetic-flow-2', 'currentValue': '3610', 'unit': 'm³/min'}]
        return body
    async with gateway(monkeypatch, response) as calls:
        first = await equipment_call(workbench, page_size=2)
        second = await equipment_call(workbench, page_num=2, page_size=2)
    assert calls[0]['path'] == '/zdsb/sbnode/selectSbNodeList'
    assert calls[0]['params'] == {'pageNum': '1', 'pageSize': '2', 'deptId': '1219', 'devOfSystem': '30'}
    assert calls[1]['params']['pageNum'] == '2'
    assert first['row_count'] == 2 and first['has_more'] is True and first['next_page'] == 2
    assert second['row_count'] == 1 and second['has_more'] is False and second['next_page'] is None
    assert first['query_id'] != second['query_id']
    data = read_data(workbench, first)
    assert data['payload'] == EQUIPMENT
    assert data['payload']['rows'][0]['currentValue'] == '3800'
    assert data['payload']['rows'][0]['device']['id'] == 'synthetic-fan'
    assert data['payload']['rows'][1]['unit'] == '℃'


@pytest.mark.asyncio
async def test_equipment_curve_uses_own_namespace_singular_node_and_raw_shape(workbench, monkeypatch):
    async with gateway(monkeypatch, CURVE) as calls:
        result = await query_fan_equipment_timeseries(workbench.identifier, runtime=workbench.runtime, **CURVE_ARGS)
    assert result['ok'], result
    assert calls[0]['path'] == '/zdsb/sbnode/selectSbNodeCurve'
    assert calls[0]['params'] == {'deptId': '1237', 'nodeId': NODE,
        'startTime': CURVE_ARGS['start_time'], 'endTime': CURVE_ARGS['end_time']}
    assert result['node_id'] == NODE and result['kind'] == 'equipment_timeseries'
    assert result['data_type'] == 'dict' and result['data_item_count'] is None
    assert result['sample_count'] is None and result['normalization'] == 'unrecognized'
    assert read_data(workbench, result)['payload'] == CURVE
    assert result['sha256'] == digest(workspace_file(workbench.user, workbench.session, result['data_path']))


@pytest.mark.asyncio
async def test_unknown_list_is_saved_without_claiming_zero_or_complete(workbench, monkeypatch):
    body = {'code': 200, 'data': {'unknownEquipmentContainer': [{'strangeField': 'value'}]}}
    async with gateway(monkeypatch, body):
        result = await equipment_call(workbench)
    assert result['ok'] and result['list_recognized'] is False
    assert result['row_count'] is None and result['has_more'] is None and result['total'] is None
    assert read_data(workbench, result)['payload'] == body


@pytest.mark.asyncio
async def test_missing_total_does_not_claim_last_page(workbench, monkeypatch):
    async with gateway(monkeypatch, {'code': 200, 'data': [{'id': 'one'}]}):
        result = await sensor_call(workbench, page_size=1)
    assert result['row_count'] == 1 and result['total'] is None and result['has_more'] is None
    assert result['next_page'] == 2


@pytest.mark.asyncio
async def test_empty_inconsistent_and_failure_are_distinct(workbench, monkeypatch):
    async with gateway(monkeypatch, {'code': 200, 'rows': [], 'total': 0}):
        empty = await equipment_call(workbench)
    assert empty['ok'] and empty['row_count'] == 0 and empty['has_more'] is False
    async with gateway(monkeypatch, {'code': 200, 'rows': [], 'total': 3}):
        inconsistent = await equipment_call(workbench)
    assert inconsistent['has_more'] is True and inconsistent['next_page'] is None
    assert '本页为空' in inconsistent['next_action']
    async with gateway(monkeypatch, proxy_status=503):
        failed = await sensor_call(workbench)
    assert failed['ok'] is False and failed['retryable'] is True and 'row_count' not in failed


@pytest.mark.asyncio
@pytest.mark.parametrize('overrides', [{'dept_id': ''}, {'page_num': 0}, {'page_size': 0},
    {'page_size': 1001}, {'sensor_type_list': ['']}, {'sensor_type_list': ['../bad']}])
async def test_invalid_list_arguments_do_not_send_a_request(workbench, monkeypatch, overrides):
    async with gateway(monkeypatch) as calls:
        result = await sensor_call(workbench, **overrides)
    assert result['ok'] is False and calls == []


@pytest.mark.asyncio
async def test_invalid_rows_do_not_become_an_empty_result(workbench, monkeypatch):
    async with gateway(monkeypatch, {'code': 200, 'rows': None, 'total': 0}):
        result = await equipment_call(workbench)
    assert result['ok'] is False and 'row_count' not in result


@pytest.mark.asyncio
async def test_large_preview_is_omitted_but_complete_data_is_kept(workbench, monkeypatch):
    body = {'code': 200, 'rows': [{'id': 'one', 'nested': 'x' * 20000}], 'total': 1}
    async with gateway(monkeypatch, body):
        result = await equipment_call(workbench)
    assert result['preview_omitted'] and result['preview'] == []
    assert read_data(workbench, result)['payload'] == body


@pytest.mark.asyncio
async def test_query_files_are_protected_and_workflow_ownership_precedes_network(workbench, monkeypatch):
    async with gateway(monkeypatch) as calls:
        wrong = SimpleNamespace(config={'configurable': {'workspace_user': 'another-user',
                                                         'workspace_session': workbench.session}})
        rejected = await sensor_call(workbench, runtime=wrong)
        assert rejected['ok'] is False and calls == []
        result = await sensor_call(workbench)
        workflow = exploration._load(workbench.identifier, workbench.user, workbench.session)
        assert exploration.manifest(workflow, workbench.user, workbench.session)['sensor_queries'] == [result]
        workspace_file(workbench.user, workbench.session, result['data_path']).write_text('{}')
        failed = await equipment_call(workbench)
        assert failed['ok'] is False and '已被修改' in failed['message'] and len(calls) == 1


class ToolModel(GenericFakeChatModel):
    def bind_tools(self, tools, *, tool_choice=None, **kwargs):
        return self

    async def _astream(self, *args, **kwargs):
        message = next(self.messages)
        yield ChatGenerationChunk(message=AIMessageChunk(content=message.content,
            tool_calls=message.tool_calls, usage_metadata=message.usage_metadata))


def query_and_deliver(workbench):
    identifier = workbench.identifier
    yield tool_call('prepare_main_fan_task', {'target': workbench.request})
    yield tool_call('query_fan_sensor_realtime', {'workflow_id': identifier, 'dept_id': '1117',
                                                'sensor_type_list': TYPES})
    yield tool_call('query_fan_equipment_realtime', {'workflow_id': identifier, 'dept_id': '1219', 'page_size': 2})
    yield tool_call('query_fan_equipment_timeseries', {'workflow_id': identifier, **CURVE_ARGS})
    workflow = main_fan_store.workflow(identifier, workbench.user, workbench.session)
    queries = workflow['context']['sensor_queries'] + workflow['context']['equipment_queries']
    paths = [item['python_paths']['data'] for item in queries]
    sources = [{'source': item['data_path'], 'pointer': '/payload'} for item in queries]
    code = f'''
from pathlib import Path
import json
documents = [json.loads(Path(path).read_text())['payload'] for path in {paths!r}]
counts = [len(documents[0]['rows']), len(documents[1]['rows']), len(documents[2]['data']['values'])]
Path('realtime-report.md').write_text('# 合成接口测试\\n\\n已读取3份真实工具输出文件，未作跨矿风机比较。')
assessment = {{'status': 'needs_materials', 'result_code': 'UNDETERMINED',
    'summary': '合成接口取数与读取测试，尚未形成主备对照。',
    'evidence': {sources!r}, 'details': {{'counts': counts}}}}
Path('realtime-assessment.json').write_text(json.dumps(assessment, ensure_ascii=False))
'''
    yield tool_call('write_file', {'file_path': '/workspace/realtime-analysis.py', 'content': code})
    yield tool_call('run_python', {'workflow_id': identifier, 'script_path': '/workspace/realtime-analysis.py',
        'output_paths': ['/workspace/realtime-report.md', '/workspace/realtime-assessment.json']})
    workflow = main_fan_store.workflow(identifier, workbench.user, workbench.session)
    receipt = workflow['context']['python_runs'][-1]
    assert receipt['exit_code'] == 0
    yield tool_call('submit_main_fan_report', {'workflow_id': identifier, 'report_path': '/workspace/realtime-report.md',
        'calculations_path': '/workspace/realtime-assessment.json', 'python_run_ids': [receipt['run_id']]})


@pytest.mark.asyncio
async def test_real_agent_queries_all_three_tools_and_archives_citations(workbench, monkeypatch, pricing_ready):
    if not capability()['available']:
        pytest.skip('当前平台缺少执行沙箱')
    responses = {'/aqjk/node/getNodeRealtimeDetails': SENSORS, '/zdsb/sbnode/selectSbNodeList': EQUIPMENT,
                 '/zdsb/sbnode/selectSbNodeCurve': CURVE}
    model = ToolModel(messages=query_and_deliver(workbench))
    monkeypatch.setattr(agent_factory.model_factory, 'get_model', lambda: model)
    agent_registry.invalidate()
    async with gateway(monkeypatch, lambda payload: responses[payload['path']]) as calls:
        async with RuntimeEnvironment(prewarm=False, billing_retry=False, recover_interrupted=False):
            result = await execute(agent_name='main_fan', user_name=workbench.user, session_id=workbench.session,
                                   message='验证三个查询入口的合成取数和沙箱读取，不跨部门作业务比较。')
    assert result.status == 'success', result.error
    assert len(calls) == 3
    value = result.structured_output
    assert value['assessment']['agent_result']['details']['counts'] == [1, 2, 2]
    evidence = value['assessment']['evidence']
    assert len(evidence) == 3 and all(item['sha256'] for item in evidence)
    assert {item['source'] for item in evidence} <= {item['url'] for item in value['sources']}
    assert len(value['assessment']['exploration']['sensor_queries']) == 1
    assert len(value['assessment']['exploration']['equipment_queries']) == 2
    assert value['platform_payload']['judgment']['resultCode'] == 'UNDETERMINED'