vent_agent_2.0/tests/test_main_fan_timeseries.py

"""真实曲线协议适配、可信取数文件及 Agent 的查询/沙箱/归档链路。"""
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_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

NODE = '61082101917301KG1010016D09'
ARGS = {'node_id': NODE, 'dept_id': '1117', 'start_time': '2026-10-19 11:00:00',
        'end_time': '2026-10-19 11:11:00'}
TIMES = ['11:00:13', '11:00:43', '11:01:15', '11:01:45', '11:02:17', '11:02:48',
         '11:03:19', '11:03:50', '11:04:20', '11:04:49', '11:05:21', '11:05:52',
         '11:06:22', '11:06:53', '11:07:24', '11:07:54', '11:08:25', '11:08:57',
         '11:09:28', '11:09:58', '11:10:30', '11:11:00']


def curve_body():
    return {'code': 200, 'msg': '操作成功', 'data': [[
        {'SensorTypeID': '1010', 'stateName': '正常', 'RealData': 1, 'NodeID': NODE,
         'StateID': '0', 'nodePlace': '主通1一级', 'DataTime': '2026-10-19 ' + value, 'KgDesc': '开'}
        for value in TIMES], [{'SensorTypeID': '1010', 'SensorName': '主通风机',
                              'NodeID': NODE, 'NodePlace': '主通1一级'}]]}


@pytest.fixture
def workbench(monkeypatch, tmp_path):
    (tmp_path / 'skills').symlink_to(config.ROOT / 'skills', target_is_directory=True)
    (tmp_path / 'db').symlink_to(config.ROOT / 'db', 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_SERIES_MAX_HOURS': '24', 'MAIN_FAN_SERIES_MAX_SAMPLES': '200000'}.items():
        monkeypatch.setenv(key, value)
    user = 'curve-user'
    session = chat_store.create_session(user)['id']
    request = {'mine_code': '610821019173', '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, manifest=prepared, runtime=runtime,
                           workspace=session_root(user, session) / 'workspace')


@asynccontextmanager
async def gateway(monkeypatch, body=None, *, status=200):
    calls = []
    body = curve_body() if body is None else body
    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 curve-test-secret'
        calls.append(json.loads(request.content))
        return httpx.Response(200, json={'code': 200, 'success': True,
            'data': {'status_code': status, 'body': deepcopy(body)}})
    client = GatewayClient(GatewaySettings('https://gateway.invalid', 'curve-test-secret'),
                           transport=httpx.MockTransport(handler))
    monkeypatch.setattr(gateway_queries, 'get_service', lambda name: client)
    try:
        yield calls
    finally:
        await client.close()


async def call(workbench, **overrides):
    return await query_fan_timeseries(workflow_id=workbench.manifest['workflow_id'], runtime=workbench.runtime,
                                      **{**ARGS, **overrides})


@pytest.mark.asyncio
async def test_curve_protocol_timezone_full_data_and_protected_sources(workbench, monkeypatch):
    async with gateway(monkeypatch) as calls:
        result = await call(workbench, start_time='2026-10-19T03:00:00Z', end_time='2026-10-19T03:11:00Z')
    assert result['ok'], result
    assert calls == [{'method': 'GET', 'path': '/aqjk/alarmcurve/getCurveByNodeId', 'params': {
        'nodeIds': NODE, 'deptId': '1117', 'startTime': ARGS['start_time'], 'endTime': ARGS['end_time']},
        'json_body': None}]
    assert result['sample_count'] == 22 and len(result['preview']) == 4
    assert result['first_time'] == '2026-10-19T11:00:13+08:00'
    assert result['last_time'] == '2026-10-19T11:11:00+08:00'
    raw = workspace_file(workbench.user, workbench.session, result['raw_path'])
    data = workspace_file(workbench.user, workbench.session, result['data_path'])
    assert json.loads(raw.read_text()) == curve_body()
    normalized = json.loads(data.read_text())
    assert len(normalized['samples']) == 22
    assert normalized['samples'][0] == {'time': result['first_time'], 'value': 1,
        'value_label': '开', 'status_code': '0', 'status_label': '正常'}
    assert normalized['metadata'][0]['NodePlace'] == '主通1一级'
    assert result['python_paths']['data'] == str(data.relative_to(workbench.workspace))
    assert result['sha256'] == digest(data)
    assert 'curve-test-secret' not in raw.read_text() + data.read_text()
    workflow = exploration._load(workbench.manifest['workflow_id'], workbench.user, workbench.session)
    assert len(workflow['context']['query_exports']) == 3
    assert exploration.manifest(workflow, workbench.user, workbench.session)['timeseries_queries'] == [result]


@pytest.mark.asyncio
async def test_repeated_queries_keep_sources_after_input_export_and_detect_tampering(workbench, monkeypatch):
    async with gateway(monkeypatch) as calls:
        first, second = await call(workbench), await call(workbench)
        assert first['query_id'] != second['query_id']
        workflow = exploration._load(workbench.manifest['workflow_id'], workbench.user, workbench.session)
        exploration._persist(workflow, workbench.user, workbench.session, '重新导出输入', {}, export=True)
        exploration.validate_inputs(workflow, workbench.user, workbench.session)
        assert len(workflow['context']['query_exports']) == 6
        workspace_file(workbench.user, workbench.session, first['data_path']).write_text('{}')
        rejected = await call(workbench)
        assert rejected['ok'] is False and '已被修改' in rejected['message']
        assert len(calls) == 2


@pytest.mark.asyncio
@pytest.mark.parametrize('overrides', [
    {'node_id': '../escape'}, {'dept_id': ''}, {'start_time': '2026-10-19'},
    {'start_time': 'not-a-time'}, {'end_time': ARGS['start_time']},
    {'end_time': '2026-10-21 11:11:00'}, {'start_time': '2026-10-19 11:00:00.5'},
])
async def test_invalid_query_does_not_call_gateway(workbench, monkeypatch, overrides):
    async with gateway(monkeypatch) as calls:
        result = await call(workbench, **overrides)
    assert result['ok'] is False and calls == []


@pytest.mark.asyncio
@pytest.mark.parametrize('damage', ['shape', 'foreign_sample', 'foreign_metadata', 'missing_time', 'bad_time'])
async def test_unrecognized_or_foreign_data_is_rejected(workbench, monkeypatch, damage):
    body = curve_body()
    if damage == 'shape':
        body['data'] = {'items': []}
    elif damage == 'foreign_sample':
        body['data'][0][0]['NodeID'] = 'foreign-node'
    elif damage == 'foreign_metadata':
        body['data'][1][0]['NodeID'] = 'foreign-node'
    elif damage == 'missing_time':
        del body['data'][0][0]['DataTime']
    else:
        body['data'][0][0]['DataTime'] = None
    async with gateway(monkeypatch, body):
        result = await call(workbench)
    assert result['ok'] is False and 'sample_count' not in result
    workflow = exploration._load(workbench.manifest['workflow_id'], workbench.user, workbench.session)
    assert not workflow['context'].get('timeseries_queries')


@pytest.mark.asyncio
async def test_empty_data_and_upstream_failure_are_distinct(workbench, monkeypatch):
    async with gateway(monkeypatch, {'code': 200, 'data': [[], []]}):
        empty = await call(workbench)
    assert empty['ok'] and empty['sample_count'] == 0 and empty['first_time'] is None
    async with gateway(monkeypatch, status=503):
        failed = await call(workbench)
    assert failed['ok'] is False and failed['retryable'] is True
    assert 'sample_count' not in failed


@pytest.mark.asyncio
async def test_quality_issues_are_reported_without_dropping_records(workbench, monkeypatch):
    body = curve_body()
    rows = body['data'][0] = [deepcopy(body['data'][0][0]) for _ in range(3)]
    rows[1]['RealData'], rows[1]['stateName'] = None, '未知'
    rows[2]['DataTime'] = '2026-10-19 10:59:00'
    async with gateway(monkeypatch, body):
        result = await call(workbench)
    assert result['sample_count'] == 3
    assert result['quality'] == {'duplicate_time_count': 1, 'out_of_order_count': 1,
                                  'out_of_range_count': 1, 'null_value_count': 1}
    assert result['preview'][1]['value'] is None and result['preview'][1]['status_label'] == '未知'


@pytest.mark.asyncio
async def test_disabled_gateway_and_sample_limit_return_actionable_feedback(workbench, monkeypatch):
    monkeypatch.setenv('SHANXI_GATEWAY_ENABLED', '0')
    async with gateway(monkeypatch) as calls:
        disabled = await call(workbench)
    assert disabled['ok'] is False and 'SHANXI_GATEWAY_ENABLED' in disabled['message'] and not calls
    monkeypatch.setenv('SHANXI_GATEWAY_ENABLED', '1')
    monkeypatch.setenv('MAIN_FAN_SERIES_MAX_SAMPLES', '2')
    async with gateway(monkeypatch):
        oversized = await call(workbench)
    assert oversized['ok'] is False and '缩小查询时间范围' in oversized['message']


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):
    yield tool_call('prepare_main_fan_task', {'target': workbench.request})
    identifier = workbench.manifest['workflow_id']
    yield tool_call('query_fan_timeseries', {'workflow_id': identifier, **ARGS})
    workflow = main_fan_store.workflow(identifier, workbench.user, workbench.session)
    query = workflow['context']['timeseries_queries'][0]
    code = f'''
from pathlib import Path
import json
data = json.loads(Path({query['python_paths']['data']!r}).read_text())
count = len(data['samples'])
Path('curve-report.md').write_text('# 时序读取测试\\n\\n实际读取 ' + str(count) + ' 条采样。')
assessment = {{'status': 'needs_materials', 'result_code': 'UNDETERMINED',
    'summary': '仅有单个开停点，尚未形成主备风机对照。',
    'evidence': [{{'source': {query['data_path']!r}, 'pointer': '/samples'}}],
    'details': {{'sample_count': count}}}}
Path('curve-assessment.json').write_text(json.dumps(assessment, ensure_ascii=False))
'''
    yield tool_call('write_file', {'file_path': '/workspace/curve-analysis.py', 'content': code})
    yield tool_call('run_python', {'workflow_id': identifier, 'script_path': '/workspace/curve-analysis.py',
        'output_paths': ['/workspace/curve-report.md', '/workspace/curve-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/curve-report.md',
        'calculations_path': '/workspace/curve-assessment.json', 'python_run_ids': [receipt['run_id']]})


@pytest.mark.asyncio
async def test_real_agent_queries_executes_and_cites_gateway_file(workbench, monkeypatch, pricing_ready):
    if not capability()['available']:
        pytest.skip('当前平台缺少执行沙箱')
    model = ToolModel(messages=query_and_deliver(workbench))
    monkeypatch.setattr(agent_factory.model_factory, 'get_model', lambda: model)
    agent_registry.invalidate()
    async with gateway(monkeypatch) 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 len(calls) == 1 and result.status == 'success', result.error
    value = result.structured_output
    assert value['assessment']['agent_result']['details']['sample_count'] == 22
    queries = value['assessment']['exploration']['timeseries_queries']
    assert value['assessment']['evidence'][0]['sha256'] == queries[0]['sha256']
    assert any(source['url'] == queries[0]['data_path'] for source in value['sources'])
    assert value['platform_payload']['judgment']['resultCode'] == 'UNDETERMINED'