from __future__ import annotations

import pytest

from chik.types.full_block import FullBlock
from chik.types.header_block import HeaderBlock
from chik.wallet.key_val_store import KeyValStore
from tests.util.db_connection import DBConnection


class TestWalletKeyValStore:
    @pytest.mark.asyncio
    async def test_store(self, bt):
        async with DBConnection(1) as db_wrapper:
            store = await KeyValStore.create(db_wrapper)
            blocks = bt.get_consecutive_blocks(20)
            block: FullBlock = blocks[0]
            block_2: FullBlock = blocks[1]

            assert (await store.get_object("a", FullBlock)) is None
            await store.set_object("a", block)
            assert await store.get_object("a", FullBlock) == block
            await store.set_object("a", block)
            assert await store.get_object("a", FullBlock) == block
            await store.set_object("a", block_2)
            await store.set_object("a", block_2)
            assert await store.get_object("a", FullBlock) == block_2
            await store.remove_object("a")
            assert (await store.get_object("a", FullBlock)) is None

            for block in blocks:
                assert (await store.get_object(block.header_hash.hex(), FullBlock)) is None
                await store.set_object(block.header_hash.hex(), block)
                assert (await store.get_object(block.header_hash.hex(), FullBlock)) == block

            # Wrong type
            await store.set_object("a", block_2)
            with pytest.raises(Exception):
                await store.get_object("a", HeaderBlock)
