harp-python/tests/test_io.py at 57918ec29891f15a780d7eb396fefd7d7d249eed · harp-tech/harp-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
from contextlib import nullcontext

import numpy as np
import pandas as pd
import pytest
from pytest import mark

from harp.data.io import REFERENCE_EPOCH, MessageType, read, to_buffer
from tests.params import DataFileParam

testdata = [
    DataFileParam(path="data/device_0.bin", expected_rows=1),
    DataFileParam(
        path="data/device_0.bin",
        expected_rows=1,
        expected_address=1,  # actual address is 0
        expected_error=ValueError,
    ),
    DataFileParam(
        path="data/device_0.bin",
        expected_rows=1,
        expected_dtype=np.dtype("uint8"),  # actual dtype is uint16
        expected_error=ValueError,
    ),
    DataFileParam(
        path="data/device_0.bin",
        expected_rows=1,
        expected_length=2,  # actual length is 1
        expected_error=ValueError,
    ),
    DataFileParam(path="data/write_0.bin", expected_address=0, expected_rows=4),
    DataFileParam(path="data/write_0.bin", expected_address=0, expected_rows=4, keep_type=True),
    DataFileParam(path="data/device_0.bin", expected_rows=300, repeat_data=300),
    DataFileParam(path="data/device_0.bin", expected_rows=1, epoch=REFERENCE_EPOCH),
    DataFileParam(path="data/empty_0.bin", expected_rows=0, epoch=REFERENCE_EPOCH),
    DataFileParam(path="data/empty_0.bin", expected_rows=0),
]


@mark.parametrize("dataFile", testdata)
def test_read(dataFile: DataFileParam):
    context = pytest.raises if dataFile.expected_error else nullcontext
    with context(dataFile.expected_error):  # type: ignore
        file_or_buf = dataFile.path
        if dataFile.repeat_data:
            with open(file_or_buf, "rb") as f:
                file_or_buf = f.read() * dataFile.repeat_data

        data = read(
            file_or_buf,
            address=dataFile.expected_address,
            dtype=dataFile.expected_dtype,
            length=dataFile.expected_length,
            epoch=dataFile.epoch,
            keep_type=dataFile.keep_type,
        )

        assert len(data) == dataFile.expected_rows
        assert isinstance(data.index, pd.DatetimeIndex if dataFile.epoch else pd.Index)
        if dataFile.keep_type:
            assert MessageType.__name__ in data.columns and data[MessageType.__name__].dtype == "category"

        if dataFile.expected_cols:
            for col in dataFile.expected_cols:
                assert col in data.columns


writedata = [
    DataFileParam(path="data/device_0.bin", expected_rows=1, expected_address=0, keep_type=True),
]


@mark.parametrize("dataFile", writedata)
def test_write(dataFile: DataFileParam):
    if dataFile.expected_address is None:
        raise AssertionError("expected address must be defined for all write tests")

    buffer = np.fromfile(dataFile.path, np.uint8)
    data = read(
        buffer,
        address=dataFile.expected_address,
        dtype=dataFile.expected_dtype,
        length=dataFile.expected_length,
        keep_type=dataFile.keep_type,
    )
    assert len(data) == dataFile.expected_rows
    write_buffer = to_buffer(data, address=dataFile.expected_address, length=dataFile.expected_length)
    assert np.array_equal(buffer, write_buffer)