netx/tests/test_biz_state_collect_finalize.py

"""Finalize / reconnect after long-collect DB disconnect."""

from __future__ import annotations

import unittest
from unittest.mock import patch

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

from netx_api.biz_state import collect_runner as runner
from netx_api.db import Base
from netx_api.models import BizStateBatch, BizStateTask


class BizStateCollectFinalizeTests(unittest.TestCase):
    def setUp(self) -> None:
        # StaticPool: all sessions share one :memory: SQLite DB.
        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-finalize",
            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-finalize",
            task_id="t-finalize",
            status="running",
            command_count=55,
            row_count=6606,
            message="",
        )
        self.db.add(self.batch)
        self.db.commit()

    def tearDown(self) -> None:
        self.db.close()

    def test_finalize_partial_after_heavy_timeout(self) -> None:
        with patch.object(runner, "SessionLocal", self.Session):
            with patch(
                "netx_api.biz_state.compare_service.schedule_auto_compare_for_task",
            ) as sched:
                status = runner._finalize_batch_status(
                    batch_id="b-finalize",
                    task_id="t-finalize",
                    cmd_count=40,
                    total_rows=1000,
                    any_fail=True,
                    any_ok=True,
                    lane_errors=["RuntimeError: biz_state_heavy_timeout (2400s)"],
                )
        self.assertEqual(status, "partial")
        sched.assert_not_called()
        self.db.expire_all()
        batch = self.db.get(BizStateBatch, "b-finalize")
        assert batch is not None
        self.assertEqual(batch.status, "partial")
        self.assertEqual(batch.command_count, 55)
        self.assertEqual(batch.row_count, 6606)
        self.assertIn("biz_state_heavy_timeout", batch.message or "")
        self.assertIsNotNone(batch.ended_at)

    def test_finalize_partial_message_from_failed_cmds_when_lane_errors_empty(self) -> None:
        from netx_api.models import BizStateBatchCommand

        self.db.add(
            BizStateBatchCommand(
                id="c-ok",
                batch_id="b-finalize",
                profile_id="zte.arp",
                parser_id="arp",
                metric_id="arp",
                raw_command="show arp",
                parse_status="ok",
                row_count=10,
            )
        )
        self.db.add(
            BizStateBatchCommand(
                id="c-aux-fail",
                batch_id="b-finalize",
                profile_id="zte.if_intf",
                parser_id="if_intf",
                metric_id="if_intf",
                raw_command="show interface brief",
                parse_status="aux_failed",
                message="aux_for=x;parse boom",
                row_count=0,
            )
        )
        self.db.commit()
        with patch.object(runner, "SessionLocal", self.Session):
            with patch(
                "netx_api.biz_state.compare_service.schedule_auto_compare_for_task",
            ):
                status = runner._finalize_batch_status(
                    batch_id="b-finalize",
                    task_id="t-finalize",
                    cmd_count=2,
                    total_rows=10,
                    any_fail=True,
                    any_ok=True,
                    lane_errors=[],
                )
        self.assertEqual(status, "partial")
        self.db.expire_all()
        batch = self.db.get(BizStateBatch, "b-finalize")
        assert batch is not None
        self.assertEqual(batch.status, "partial")
        self.assertTrue(str(batch.message or "").strip())
        self.assertIn("aux_failed", (batch.message or "").lower())
        self.assertIn("show interface brief", batch.message or "")
        self.assertIn("parse boom", batch.message or "")

    def test_finalize_partial_when_skipped_bindings_with_success(self) -> None:
        from netx_api.models import BizStateBatchCommand

        self.db.add(
            BizStateBatchCommand(
                id="c-ok2",
                batch_id="b-finalize",
                profile_id="zte.arp",
                raw_command="show arp",
                parse_status="ok",
                row_count=1,
            )
        )
        self.db.add(
            BizStateBatchCommand(
                id="c-skip",
                batch_id="b-finalize",
                profile_id="zte.bgp_vpnv4_neighbor_in",
                raw_command="show bgp vpnv4 ...",
                parse_status="skipped",
                message="profile requires parameter bindings",
                row_count=0,
            )
        )
        self.db.commit()
        with patch.object(runner, "SessionLocal", self.Session):
            with patch(
                "netx_api.biz_state.compare_service.schedule_auto_compare_for_task",
            ):
                status = runner._finalize_batch_status(
                    batch_id="b-finalize",
                    task_id="t-finalize",
                    cmd_count=2,
                    total_rows=1,
                    any_fail=False,
                    any_ok=True,
                    lane_errors=[],
                )
        self.assertEqual(status, "partial")
        self.db.expire_all()
        batch = self.db.get(BizStateBatch, "b-finalize")
        assert batch is not None
        self.assertIn("skipped", (batch.message or "").lower())
        self.assertIn("zte.bgp_vpnv4_neighbor_in", batch.message or "")
        self.assertIn("bindings", (batch.message or "").lower())

    def test_finalize_partial_lists_unmatched_and_failed_commands(self) -> None:
        from netx_api.models import BizStateBatchCommand

        self.db.add(
            BizStateBatchCommand(
                id="c-ok3",
                batch_id="b-finalize",
                profile_id="zte.arp",
                metric_id="arp",
                raw_command="show arp",
                parse_status="ok",
                row_count=1,
            )
        )
        self.db.add(
            BizStateBatchCommand(
                id="c-unmatched",
                batch_id="b-finalize",
                profile_id="zte.legacy_x",
                raw_command="show weird-legacy",
                parse_status="unmatched",
                message="no profile matched concrete command",
                row_count=0,
            )
        )
        self.db.add(
            BizStateBatchCommand(
                id="c-fail",
                batch_id="b-finalize",
                profile_id="zte.bgp_route",
                metric_id="bgp_route",
                raw_command="show bgp vpnv4 unicast neighbor in 1.1.1.1",
                parse_status="failed",
                message="parse: ValueError: boom",
                row_count=0,
            )
        )
        self.db.commit()
        with patch.object(runner, "SessionLocal", self.Session):
            with patch(
                "netx_api.biz_state.compare_service.schedule_auto_compare_for_task",
            ):
                status = runner._finalize_batch_status(
                    batch_id="b-finalize",
                    task_id="t-finalize",
                    cmd_count=3,
                    total_rows=1,
                    any_fail=True,
                    any_ok=True,
                    lane_errors=[],
                )
        self.assertEqual(status, "partial")
        self.db.expire_all()
        batch = self.db.get(BizStateBatch, "b-finalize")
        assert batch is not None
        msg = batch.message or ""
        self.assertIn("unmatched=1", msg)
        self.assertIn("failed=1", msg)
        self.assertIn("show weird-legacy", msg)
        self.assertIn("show bgp vpnv4", msg)
        self.assertIn("zte.bgp_route", msg)
        self.assertNotIn("partial success (some steps failed", msg)

    def test_finalize_retries_once_on_operational_error(self) -> None:
        calls = {"n": 0}
        real_session = self.Session

        class FlakySession:
            """First commit raises OperationalError; subsequent sessions work."""

            def __init__(self) -> None:
                self._inner = real_session()
                self._failed = False

            def get(self, *args, **kwargs):
                return self._inner.get(*args, **kwargs)

            def query(self, *args, **kwargs):
                return self._inner.query(*args, **kwargs)

            def add(self, *args, **kwargs):
                return self._inner.add(*args, **kwargs)

            def commit(self) -> None:
                calls["n"] += 1
                if calls["n"] == 1:
                    self._failed = True
                    raise OperationalError(
                        "UPDATE",
                        {},
                        Exception(
                            "consuming input failed: server closed the connection unexpectedly"
                        ),
                    )
                self._inner.commit()

            def rollback(self) -> None:
                try:
                    self._inner.rollback()
                except Exception:
                    pass

            def close(self) -> None:
                self._inner.close()

            def connection(self):
                # Avoid invalidate() wiping StaticPool's only connection in tests.
                raise RuntimeError("skip invalidate in test")

            def add(self, *args, **kwargs):
                return self._inner.add(*args, **kwargs)

        def session_factory():
            return FlakySession()

        with patch.object(runner, "SessionLocal", session_factory):
            with patch(
                "netx_api.biz_state.compare_service.schedule_auto_compare_for_task",
            ):
                status = runner._finalize_batch_status(
                    batch_id="b-finalize",
                    task_id="t-finalize",
                    cmd_count=55,
                    total_rows=6606,
                    any_fail=True,
                    any_ok=True,
                    lane_errors=["RuntimeError: biz_state_heavy_timeout (2400s)"],
                )
        self.assertEqual(status, "partial")
        self.assertEqual(calls["n"], 2)
        self.db.expire_all()
        batch = self.db.get(BizStateBatch, "b-finalize")
        assert batch is not None
        self.assertEqual(batch.status, "partial")
        self.assertIn("biz_state_heavy_timeout", batch.message or "")

    def test_finalize_success_schedules_auto_compare(self) -> None:
        with patch.object(runner, "SessionLocal", self.Session):
            with patch(
                "netx_api.biz_state.compare_service.schedule_auto_compare_for_task",
            ) as sched:
                status = runner._finalize_batch_status(
                    batch_id="b-finalize",
                    task_id="t-finalize",
                    cmd_count=55,
                    total_rows=6606,
                    any_fail=False,
                    any_ok=True,
                    lane_errors=[],
                )
        self.assertEqual(status, "success")
        sched.assert_called_once_with("t-finalize", "b-finalize")

    def test_fail_batch_retries_on_stale_connection(self) -> None:
        calls = {"n": 0}
        real_session = self.Session

        class FlakySession:
            def __init__(self) -> None:
                self._inner = real_session()

            def get(self, *args, **kwargs):
                return self._inner.get(*args, **kwargs)

            def commit(self) -> None:
                calls["n"] += 1
                if calls["n"] == 1:
                    raise OperationalError(
                        "UPDATE",
                        {},
                        Exception("server closed the connection unexpectedly"),
                    )
                self._inner.commit()

            def rollback(self) -> None:
                try:
                    self._inner.rollback()
                except Exception:
                    pass

            def close(self) -> None:
                self._inner.close()

            def connection(self):
                raise RuntimeError("skip invalidate in test")

        with patch.object(runner, "SessionLocal", FlakySession):
            runner._fail_batch_status("b-finalize", "RuntimeError: boom")
        self.assertEqual(calls["n"], 2)
        self.db.expire_all()
        batch = self.db.get(BizStateBatch, "b-finalize")
        assert batch is not None
        self.assertEqual(batch.status, "failed")
        self.assertIn("boom", batch.message or "")

    def test_run_db_with_reconnect_reraises_non_stale(self) -> None:
        calls = {"n": 0}

        class BoomSession:
            def commit(self) -> None:
                pass

            def rollback(self) -> None:
                pass

            def close(self) -> None:
                pass

            def connection(self):
                raise RuntimeError("no conn")

        def fn(db) -> None:
            calls["n"] += 1
            raise ValueError("not a disconnect")

        with patch.object(runner, "SessionLocal", BoomSession):
            with self.assertRaises(ValueError):
                runner._run_db_with_reconnect(fn, label="test")
        self.assertEqual(calls["n"], 1)

    def test_is_stale_db_connection(self) -> None:
        self.assertTrue(
            runner._is_stale_db_connection(
                OperationalError("x", {}, Exception("server closed the connection"))
            )
        )
        self.assertTrue(
            runner._is_stale_db_connection(
                RuntimeError("consuming input failed: server closed the connection unexpectedly")
            )
        )
        self.assertFalse(runner._is_stale_db_connection(ValueError("nope")))


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