caiwuzongzhang/tests/ledger_helpers.py

"""Shared fixtures for B-44 ledger / manual / position tests."""

from __future__ import annotations

from pathlib import Path
import sqlite3
import tempfile
import unittest

from bank_importer import auth, matching, master_data
from bank_importer.db import connect, migrate, utc_now


class LedgerBase(unittest.TestCase):
    def setUp(self) -> None:
        self.temp_dir = tempfile.TemporaryDirectory()
        self.addCleanup(self.temp_dir.cleanup)
        root = Path(self.temp_dir.name)
        self.db_path = root / "app.db"
        self.connection = connect(self.db_path)
        self.addCleanup(self.connection.close)
        migrate(self.connection)
        self.admin = self._admin()
        self.company_a = self._company("甲公司")
        self.company_b = self._company("乙公司")
        self.company_c = self._company("丙公司")
        self.account_a = self._approved_account(self.company_a, "6222000000000001")
        self.account_b = self._approved_account(self.company_b, "6222000000000002")

    def _admin(self):
        auth.create_user(self.connection, "admin-u", "AdminPass123", "admin")
        return self.connection.execute(
            "SELECT * FROM users WHERE username = 'admin-u'"
        ).fetchone()

    def _company(self, name: str) -> int:
        with self.connection:
            cursor = self.connection.execute(
                "INSERT INTO companies (name, created_at, updated_at) VALUES (?, ?, ?)",
                (name, utc_now(), utc_now()),
            )
        return int(cursor.lastrowid)

    def _approved_account(self, company_id: int, number: str, start: str = "2026-01-01"):
        account = master_data.submit_bank_account(
            self.connection, company_id=company_id, bank_name="中信银行",
            account_type="基本户", account_number=number, start_date=start,
            actor=None,
        )
        return master_data.review_bank_account(
            self.connection, account["id"], "approve", None, self.admin,
            effective_from=start,
        )

    def add_row(
        self,
        company_id: int,
        *,
        own_account: str,
        cp_account: str | None = None,
        income: str = "0",
        expense: str = "0",
        at: str = "2026-01-05T10:00:00",
        currency: str = "CNY",
        reference: str | None = None,
        summary: str | None = None,
        purpose: str | None = None,
    ) -> int:
        with self.connection:
            cursor = self.connection.execute(
                """
                INSERT INTO source_files (sha256, original_filename, size_bytes, storage_path, created_at)
                VALUES (?, '测试.xlsx', 1, 'data/files/测试.xlsx', ?)
                """,
                (utc_now(), utc_now()),
            )
            source_file_id = int(cursor.lastrowid)
            cursor = self.connection.execute(
                """
                INSERT INTO import_batches (source_file_id, status, company_id, created_at, updated_at)
                VALUES (?, 'parsing', ?, ?, ?)
                """,
                (source_file_id, company_id, utc_now(), utc_now()),
            )
            batch_id = int(cursor.lastrowid)
            cursor = self.connection.execute(
                """
                INSERT INTO sheet_batches (
                    import_batch_id, sheet_name, bank_name, template_id, template_version,
                    header_row, own_account, own_name, period_start, period_end,
                    transaction_count, warnings, created_at
                ) VALUES (?, '流水', '测试银行', 'test-v1', 1, 1, NULL, NULL, NULL, NULL, 1, '[]', ?)
                """,
                (batch_id, utc_now()),
            )
            sheet_batch_id = int(cursor.lastrowid)
            self.connection.execute(
                """
                INSERT INTO sheet_reviews (
                    import_batch_id, sheet_name, outcome, sheet_batch_id,
                    review_status, created_at
                ) VALUES (?, '流水', 'parsed', ?, 'confirmed', ?)
                """,
                (batch_id, sheet_batch_id, utc_now()),
            )
            cursor = self.connection.execute(
                """
                INSERT INTO source_rows (
                    sheet_batch_id, source_row, transaction_at, income, expense, balance,
                    own_account, own_name, counterparty_account, counterparty_name,
                    counterparty_bank, summary, purpose, reference, currency, created_at
                ) VALUES (?, 1, ?, ?, ?, NULL, ?, '测试', ?, '对方', NULL, ?, ?, ?, ?, ?)
                """,
                (
                    sheet_batch_id, at, income, expense, own_account,
                    cp_account, summary, purpose, reference, currency, utc_now(),
                ),
            )
        return int(cursor.lastrowid)

    def pair(
        self,
        payer: int,
        payee: int,
        amount: str,
        at: str = "2026-01-05T10:00:00",
        *,
        currency: str = "CNY",
        summary: str = "借款",
        purpose: str = "往来款",
    ) -> tuple[int, int]:
        """Create a mirrored A/B pair and reconcile into an eligible event."""
        payer_account = self._account_of(payer)
        payee_account = self._account_of(payee)
        row_payer = self.add_row(
            payer, own_account=payer_account, cp_account=payee_account,
            expense=amount, at=at, currency=currency, summary=summary, purpose=purpose,
        )
        row_payee = self.add_row(
            payee, own_account=payee_account, cp_account=payer_account,
            income=amount, at=at.replace("T10:", "T11:"), currency=currency,
            summary=summary, purpose=purpose,
        )
        matching.reconcile_rows(self.connection, [row_payer, row_payee])
        return row_payer, row_payee

    def _account_of(self, company_id: int) -> str:
        if company_id == self.company_a:
            return self.account_a["account_number"]
        if company_id == self.company_b:
            return self.account_b["account_number"]
        return "6222000000000005"

    def eligible(self) -> list[sqlite3.Row]:
        return matching.eligible_intercompany_events(self.connection)

    def ledger_events(self) -> list[sqlite3.Row]:
        return self.connection.execute(
            "SELECT * FROM ledger_events ORDER BY id"
        ).fetchall()

    def current(self, ledger_event_id: int) -> sqlite3.Row | None:
        return self.connection.execute(
            """
            SELECT r.* FROM current_ledger_event_revisions c
            JOIN ledger_event_revisions r ON r.id = c.revision_id
            WHERE c.ledger_event_id = ?
            """,
            (ledger_event_id,),
        ).fetchone()

    def position_events(self) -> list[sqlite3.Row]:
        return self.connection.execute(
            "SELECT * FROM eligible_position_events ORDER BY ledger_event_id"
        ).fetchall()