netx/tests/test_biz_state_spool.py

"""biz_state collect spool: disk collect + batched DB flush."""

from __future__ import annotations

import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
from uuid import uuid4

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import StaticPool

from netx_api.biz_state import collect_runner as runner
from netx_api.biz_state import spool as spool_mod
from netx_api.biz_state.spool import (
    SpooledCommand,
    clear_batch_spool,
    read_raw_text,
    read_records,
    write_raw_text,
    write_records,
)
from netx_api.db import Base
from netx_api.models import BizStateBatch, BizStateBatchCommand, BizStateMetricRow, BizStateTask


class BizStateSpoolIoTests(unittest.TestCase):
    def setUp(self) -> None:
        self._tmpdir = tempfile.TemporaryDirectory()
        self.root = Path(self._tmpdir.name)
        self._patcher = patch.object(spool_mod.settings, "biz_state_spool_dir", str(self.root))
        self._patcher.start()

    def tearDown(self) -> None:
        self._patcher.stop()
        self._tmpdir.cleanup()

    def test_write_read_raw_and_records(self) -> None:
        bid = "batch1"
        cid = "cmd1"
        rel = write_raw_text(bid, cid, "show arp\nA B C")
        self.assertTrue(rel.endswith("cmd1.raw.txt"))
        self.assertEqual(read_raw_text(rel), "show arp\nA B C")
        rrel = write_records(bid, cid, [{"ip": "1.1.1.1"}, {"ip": "2.2.2.2"}])
        recs = read_records(rrel)
        self.assertEqual(len(recs), 2)
        self.assertEqual(recs[0]["ip"], "1.1.1.1")

    def test_raw_max_bytes_truncate(self) -> None:
        bid = "b2"
        cid = "c2"
        rel = write_raw_text(bid, cid, "x" * 100)
        text = read_raw_text(rel, max_bytes=20)
        self.assertIn("truncated", text)
        self.assertLess(len(text), 80)

    def test_clear_batch_spool(self) -> None:
        bid = "b3"
        write_raw_text(bid, "c", "hi")
        self.assertTrue((self.root / "b3").is_dir())
        clear_batch_spool(bid)
        self.assertFalse((self.root / "b3").exists())


class BizStateFlushSpoolTests(unittest.TestCase):
    def setUp(self) -> None:
        self._tmpdir = tempfile.TemporaryDirectory()
        self.root = Path(self._tmpdir.name)
        self._spool_patch = patch.object(
            spool_mod.settings, "biz_state_spool_dir", str(self.root)
        )
        self._spool_patch.start()

        engine = create_engine(
            "sqlite+pysqlite:///:memory:",
            future=True,
            connect_args={"check_same_thread": False},
            poolclass=StaticPool,
        )
        TestingSession = sessionmaker(
            bind=engine, autoflush=False, autocommit=False, expire_on_commit=False
        )
        Base.metadata.create_all(bind=engine)
        self.Session = TestingSession
        self.db = TestingSession()
        self.task = BizStateTask(
            id="t-spool",
            source="managed",
            ne_id="ne1",
            ne_name="PE1",
            status="running",
            collect_running=True,
            interval_sec=300,
        )
        self.db.add(self.task)
        self.batch = BizStateBatch(
            id="b-spool",
            task_id="t-spool",
            status="running",
            command_count=0,
            row_count=0,
        )
        self.db.add(self.batch)
        self.db.commit()
        self._session_patch = patch.object(runner, "SessionLocal", TestingSession)
        self._session_patch.start()

    def tearDown(self) -> None:
        self._session_patch.stop()
        self._spool_patch.stop()
        self.db.close()
        self._tmpdir.cleanup()

    def test_flush_inserts_command_and_metric_rows(self) -> None:
        cid = uuid4().hex
        raw_rel = write_raw_text("b-spool", cid, "ARP OUTPUT")
        rec_rel = write_records(
            "b-spool",
            cid,
            [{"ip": "10.0.0.1", "mac": "aaaa"}, {"ip": "10.0.0.2", "mac": "bbbb"}],
        )
        pending = [
            SpooledCommand(
                id=cid,
                batch_id="b-spool",
                task_item_id="item1",
                profile_id="zte.arp",
                parser_id="zte_arp",
                metric_id="arp",
                raw_command="show arp",
                parse_status="ok",
                message="spooled",
                raw_rel_path=raw_rel,
                records_rel_path=rec_rel,
                persist_kind="metric",
            )
        ]
        cmds, rows = runner._flush_spooled_commands("b-spool", pending)
        self.assertEqual(cmds, 1)
        self.assertEqual(rows, 2)
        self.assertEqual(pending, [])
        self.db.expire_all()
        cmd = self.db.get(BizStateBatchCommand, cid)
        assert cmd is not None
        self.assertEqual(cmd.parse_status, "ok")
        self.assertEqual(cmd.raw_text, "ARP OUTPUT")
        self.assertEqual(cmd.row_count, 2)
        n = (
            self.db.query(BizStateMetricRow)
            .filter(BizStateMetricRow.batch_command_id == cid)
            .count()
        )
        self.assertEqual(n, 2)
        batch = self.db.get(BizStateBatch, "b-spool")
        assert batch is not None
        self.assertEqual(batch.command_count, 1)
        self.assertEqual(batch.row_count, 2)

    def test_flush_batches_multiple_without_per_cmd_sessions(self) -> None:
        pending: list[SpooledCommand] = []
        for i in range(5):
            cid = uuid4().hex
            raw_rel = write_raw_text("b-spool", cid, f"out-{i}")
            pending.append(
                SpooledCommand(
                    id=cid,
                    batch_id="b-spool",
                    raw_command=f"show x {i}",
                    parse_status="skipped_custom",
                    message="custom_raw",
                    raw_rel_path=raw_rel,
                )
            )
        cmds, rows = runner._flush_spooled_commands("b-spool", pending)
        self.assertEqual(cmds, 5)
        self.assertEqual(rows, 0)
        self.db.expire_all()
        n = (
            self.db.query(BizStateBatchCommand)
            .filter(BizStateBatchCommand.batch_id == "b-spool")
            .count()
        )
        self.assertEqual(n, 5)


if __name__ == "__main__":
    unittest.main()