예제 소스

test_evaluate_qwen_chart_lora.py

import math
import tempfile
import unittest
from pathlib import Path

from evaluate_qwen_chart_lora import (
    FEE,
    PaperAccount,
    _duration_expired,
    _truncate_predictions,
    build_prompt,
    parse_action,
)


class ParseActionTest(unittest.TestCase):
    def test_exact_actions(self):
        self.assertEqual(parse_action("매수"), (1, True))
        self.assertEqual(parse_action("매도"), (0, True))
        self.assertEqual(parse_action("유지"), (2, True))

    def test_embedded_action_and_invalid_fallback(self):
        self.assertEqual(parse_action("판단: 매도\n"), (0, True))
        self.assertEqual(parse_action("알 수 없음"), (2, False))


class PaperAccountTest(unittest.TestCase):
    def test_account_buy_and_sell_match_market_rules(self):
        account = PaperAccount.create(10_000)

        self.assertTrue(account.execute(1, 1_000))
        expected_unit = (9_950 / 1_000) * (1.0 - FEE)
        self.assertTrue(math.isclose(account.balance, 50.0))
        self.assertTrue(math.isclose(account.unit, expected_unit))
        self.assertFalse(account.execute(1, 1_000), "보유 중 추가 매수는 무시")

        self.assertTrue(account.execute(0, 1_100))
        expected_balance = 50.0 + expected_unit * 1_100 * (1.0 - FEE)
        self.assertTrue(math.isclose(account.balance, expected_balance))
        self.assertEqual(account.unit, 0.0)
        self.assertEqual(account.completed_trades, 1)
        self.assertEqual(account.profitable_trades, 1)
        self.assertFalse(account.execute(0, 1_100), "미보유 매도는 무시")

    def test_minimum_order_is_enforced(self):
        account = PaperAccount.create(5_000)
        self.assertFalse(account.execute(1, 1_000))
        self.assertEqual(account.balance, 5_000)

    def test_mark_to_market_and_liquidation_value(self):
        account = PaperAccount.create(10_000)
        account.execute(1, 1_000)
        self.assertGreater(account.portfolio_value(1_100), account.liquidation_value(1_100))


class PromptTest(unittest.TestCase):
    def test_prompt_records_position_state(self):
        self.assertIn("ETH 보유 없음", build_prompt(False, 0, 1_000))
        prompt = build_prompt(True, 1_000, 1_100)
        self.assertIn("ETH 보유 중", prompt)
        self.assertIn("진입가 1,000원", prompt)


class ResumeCsvTest(unittest.TestCase):
    def test_truncate_predictions_removes_rows_after_checkpoint(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "predictions.csv"
            path.write_text(
                "bar_index,action\n300,1\n301,2\n302,0\n",
                encoding="utf-8",
            )
            _truncate_predictions(path, 301)
            self.assertEqual(
                path.read_text(encoding="utf-8").splitlines(),
                ["bar_index,action", "300,1", "301,2"],
            )


class RandomStartArgsTest(unittest.TestCase):
    def test_random_start_flag_exists(self):
        from evaluate_qwen_chart_lora import DEFAULT_START_BAR

        self.assertEqual(DEFAULT_START_BAR, 300)


class DurationLimitTest(unittest.TestCase):
    def test_duration_limit(self):
        self.assertFalse(_duration_expired(100.0, None, 1_000.0))
        self.assertFalse(_duration_expired(100.0, 3_600.0, 3_699.9))
        self.assertTrue(_duration_expired(100.0, 3_600.0, 3_700.0))


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