autoChatGPTPlus-Paypal/test_storage.py

import os
import tempfile
import unittest
from unittest.mock import patch

import storage


class TrialAccountStorageTests(unittest.TestCase):
    def setUp(self):
        self.tempdir = tempfile.TemporaryDirectory()
        self.addCleanup(self.tempdir.cleanup)

        self.data_dir_patch = patch.object(storage, "DATA_DIR", self.tempdir.name)
        self.db_path_patch = patch.object(storage, "DB_PATH", os.path.join(self.tempdir.name, "accounts.db"))
        self.data_dir_patch.start()
        self.db_path_patch.start()
        self.addCleanup(self.data_dir_patch.stop)
        self.addCleanup(self.db_path_patch.stop)

        storage.init_db()

    def test_trial_eligibility_remains_after_final_status_changes(self):
        storage.upsert_account(
            "trial@example.com",
            "pw",
            fields={
                "final_status": "paid",
                "initial_session": {
                    "accessToken": "tok-trial",
                    "account": {"planType": "free"},
                },
                "notes": storage.TRIAL_ELIGIBLE_NOTE,
            },
        )
        storage.upsert_account(
            "trial@example.com",
            "pw",
            fields={
                "final_status": "cpa_uploaded",
                "plus_session": {
                    "accessToken": "tok-trial",
                    "account": {"planType": "plus"},
                },
            },
        )

        account = storage.get_account("trial@example.com")

        self.assertTrue(account["is_trial_account"])
        self.assertEqual(account["trial_state"], "eligible")
        self.assertTrue(account["can_retry_payment"])

    def test_get_account_keeps_manual_payment_trial_separate_from_trial_account(self):
        storage.upsert_account(
            "manual@example.com",
            "pw",
            fields={
                "final_status": "trial",
                "last_error": "非免费金额需手动付款: 金额为 $20.00",
                "initial_session": {
                    "accessToken": "tok-manual",
                    "account": {"planType": "free"},
                },
                "notes": storage.TRIAL_INELIGIBLE_NOTE,
            },
        )

        account = storage.get_account("manual@example.com")

        self.assertFalse(account["is_trial_account"])
        self.assertEqual(account["trial_state"], "ineligible")
        self.assertTrue(account["can_retry_payment"])

    def test_list_accounts_trial_filter_includes_trial_accounts_and_pending_payment_accounts(self):
        storage.upsert_account(
            "trial@example.com",
            "pw",
            fields={
                "final_status": "cpa_uploaded",
                "initial_session": {
                    "accessToken": "tok-trial",
                    "account": {"planType": "free"},
                },
                "notes": storage.TRIAL_ELIGIBLE_NOTE,
            },
        )
        storage.upsert_account(
            "manual@example.com",
            "pw",
            fields={
                "final_status": "trial",
                "last_error": "非免费金额需手动付款: 金额为 $20.00",
                "initial_session": {
                    "accessToken": "tok-manual",
                    "account": {"planType": "free"},
                },
                "notes": storage.TRIAL_INELIGIBLE_NOTE,
            },
        )
        storage.upsert_account(
            "plain@example.com",
            "pw",
            fields={
                "final_status": "cpa_uploaded",
                "initial_session": {
                    "accessToken": "tok-plain",
                    "account": {"planType": "free"},
                },
            },
        )

        accounts = storage.list_accounts(status="trial")

        self.assertEqual(
            [account["email"] for account in accounts],
            ["manual@example.com", "trial@example.com"],
        )


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