Repository: Chia-Network/chia-blockchain
File: chia/_tests/core/data_layer/test_data_store.py (Python)

Code before:
    retrieved_keys_values = {node.key: node.value for node in terminal_nodes}

    assert retrieved_keys_values == inserted_keys_values


@pytest.mark.anyio
@boolean_datacases(name="test_confirmation", true="confirm blobs", false="don't confirm blobs")
async def test_unconfirmed_keys_values(
    data_store: DataStore,
    store_id: bytes32,
    seeded_random: random.Random,
    test_confirmation: bool,
) -> None:
    num_keys = 10000
    num_confirmed_files = 0
    num_files = 0
    num_confirmed_blobs = 0

    async with data_store.db_wrapper.writer() as writer:
        for _ in range(num_keys):
            confirmed = seeded_random.choice([True, False])
            use_file = seeded_random.choice([True, False])
            assert data_store.prefer_db_kv_blob_length > 7
            size = data_store.prefer_db_kv_blob_length + 1 if use_file else 8
            key = bytes(seeded_random.getrandbits(8) for _ in range(size))
            value = bytes(seeded_random.getrandbits(8) for _ in range(size))
            await data_store.add_key_value(key, value, store_id, writer, confirmed)
            if use_file:
                num_files += 1
                if confirmed:
                    num_confirmed_files += confirmed
            num_confirmed_blobs += confirmed

    async with data_store.db_wrapper.reader() as reader:
        async with reader.execute("SELECT COUNT(*) FROM ids WHERE confirmed = 0") as cursor:
            row_count = await cursor.fetchone()
            assert row_count is not None
            assert row_count[0] == 2 * (num_keys - num_confirmed_blobs)

    keys_value_path = data_store.key_value_blobs_path.joinpath(store_id.hex())
    assert sum(1 for _ in keys_value_path.rglob("*") if _.is_file()) == 2 * num_files

    if test_confirmation:
        await data_store.confirm_all_kvids(store_id)
    expected_num_files = 2 * num_files if test_confirmation else 2 * num_confirmed_files

    await data_store.delete_unconfirmed_kvids(store_id)
    assert sum(1 for _ in keys_value_path.rglob("*") if _.is_file()) == expected_num_files

    async with data_store.db_wrapper.reader() as reader:
        async with reader.execute("SELECT COUNT(*) FROM ids WHERE confirmed = 0") as cursor:

Review comment (suggestion):
https://docs.python.org/3.13/library/random.html#random.randbytes
> _Added in version 3.9._

we're finally allowed to use `.randbytes()`!  `:]`  :tada: 

```suggestion
            key = seeded_random.randbytes(size)
            value = seeded_random.randbytes(size)
```

this reminded me to make https://github.com/Chia-Network/chia-blockchain/pull/19838

Code after:

    terminal_nodes = await data_store.get_keys_values(store_id=store_id)
    retrieved_keys_values = {node.key: node.value for node in terminal_nodes}

    assert retrieved_keys_values == inserted_keys_values


@pytest.mark.anyio
@boolean_datacases(name="success", false="invalid file", true="valid file")
async def test_db_data_insert_from_file(
    data_store: DataStore,
    store_id: bytes32,
    tmp_path: Path,
    seeded_random: random.Random,
    success: bool,
) -> None:
    num_keys = 1000
    db_uri = generate_in_memory_db_uri()

    async with DataStore.managed(
        database=db_uri,
        uri=True,
        merkle_blobs_path=tmp_path.joinpath("merkle-blobs-tmp"),
        key_value_blobs_path=tmp_path.joinpath("key-value-blobs-tmp"),
    ) as tmp_data_store:
        await tmp_data_store.create_tree(store_id, status=Status.COMMITTED)
        changelist: list[dict[str, Any]] = []
        for _ in range(num_keys):
            use_file = seeded_random.choice([True, False])
            assert tmp_data_store.prefer_db_kv_blob_length > 7
            size = tmp_data_store.prefer_db_kv_blob_length + 1 if use_file else 8
            key = seeded_random.randbytes(size)
            value = seeded_random.randbytes(size)
            changelist.append({"action": "insert", "key": key, "value": value})

        await tmp_data_store.insert_batch(store_id, changelist, status=Status.COMMITTED)
        root = await tmp_data_store.get_tree_root(store_id)
        files_path = tmp_path.joinpath("files")
        await write_files_for_root(tmp_data_store, store_id, root, files_path, 1000)
        assert root.node_hash is not None
        filename = get_delta_filename_path(files_path, store_id, root.node_hash, 1)
        assert filename.exists()

    root_hash = bytes32([0] * 31 + [1]) if not success else root.node_hash
    sinfo = ServerInfo("http://127.0.0.1/8003", 0, 0)

    if not success:
        target_filename_path = get_delta_filename_path(files_path, store_id, root_hash, 1)
        shutil.copyfile(filename, target_filename_path)
        assert target_filename_path.exists()