kraken/backend/tests/device/sig_chain/test_sig_reader.py

"""单元测试 sig_reader"""
import collections
import os

from device.sig_chain.sig_reader import Reader

TEST_DATA_PATH = "tests/data/"
BDF_FILE_PATH = os.path.join(TEST_DATA_PATH, "5_3_right_hand.bdf")


def test_read():
    ch_names = [
        "Fz", "Fp1", "F3", "F7", "C3", "T3", "T5", "P3", "O1", "Cz", "Oz", "Pz",
        "O2", "P4", "T6", "T4", "C4", "F8", "F4", "Fp2"
    ]
    reader = Reader()
    raw = reader.read(BDF_FILE_PATH, tuple(ch_names))
    assert (20, 386000) == raw.get_data().shape


def test_fix_annotation():
    ch_names = [
        "Fz", "Fp1", "F3", "F7", "C3", "T3", "T5", "P3", "O1", "Cz", "Oz", "Pz",
        "O2", "P4", "T6", "T4", "C4", "F8", "F4", "Fp2"
    ]
    reader = Reader()
    raw = reader.read(BDF_FILE_PATH, tuple(ch_names))
    reader.fix_annotation(raw)

    ret = collections.Counter(raw.annotations.description)
    assert 1 == ret["initialRest"]
    assert 15 == ret["mi"]
    assert 15 == ret["rest"]