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

"""单元测试 sig_save"""
import os

import numpy as np

from device.sig_chain.sig_save import SigSave
from schemas.subjects import SubjectCreate

channel_labels = [
    'T6', 'P4', 'Pz', 'M2', 'F8', 'F4', 'Fp1', 'Cz', 'M1', 'F7', 'F3', 'C3',
    'T3', 'A1', 'Oz', 'O1', 'O2', 'Fz', 'C4', 'T4', 'Fp2', 'A2', 'T5', 'P3'
]
test_data_path = './tests/core/sig_chain/test_data/'

filename = 'testfilename.bdf'

TASK_PER_RUN = 1


def setup_module():
    if not os.path.exists(test_data_path):
        os.makedirs(test_data_path)


def teardown_module():
    os.removedirs(test_data_path)


def create_subject():
    return SubjectCreate(name='nobody',
                         id_card='12345',
                         gender='男',
                         birthday='1988-01-01',
                         rehabilitation_parts=['左手'])


def test_subject_set_edf_header_success():
    saver = SigSave(channel_labels, 1000, 375000, -375000)
    subject = create_subject()
    saver.set_edf_header(subject, filename, TASK_PER_RUN, test_data_path)
    assert saver.is_ready is True


def test_close_edf_file():
    saver = SigSave(channel_labels, 1000, 375000, -375000)
    subject = create_subject()
    saver.set_edf_header(subject, filename, TASK_PER_RUN, test_data_path)
    saver.close_edf_file()
    assert saver.is_ready is False


def test_save_raw_data_once():
    saver = SigSave(channel_labels, 1000, 375000, -375000)
    subject = create_subject()
    saver.set_edf_header(subject, filename, TASK_PER_RUN, test_data_path)
    data = np.ones((len(channel_labels), 1000))
    saver.save_raw_data(data)
    file_path = test_data_path + filename
    assert os.path.exists(file_path)
    assert os.path.getsize(file_path) > 0
    saver.close_edf_file()
    os.remove(file_path)


def test_save_raw_data_10_times():
    saver = SigSave(channel_labels, 1000, 375000, -375000)
    subject = create_subject()
    saver.set_edf_header(subject, filename, TASK_PER_RUN, test_data_path)
    data = np.ones((len(channel_labels), 1000))
    for _ in range(10):
        saver.save_raw_data(data)
    file_path = test_data_path + filename
    assert os.path.exists(file_path)
    assert os.path.getsize(file_path) > 0
    saver.close_edf_file()
    os.remove(file_path)


def test_save_raw_data_without_set_header():
    saver = SigSave(channel_labels, 1000, 375000, -375000)
    data = np.ones((len(channel_labels), 1000))
    saver.save_raw_data(data)
    file_path = test_data_path + filename
    assert os.path.exists(file_path) is False


def test_edf_data_mark():
    saver = SigSave(channel_labels, 1000, 375000, -375000)
    subject = create_subject()
    saver.set_edf_header(subject, filename, TASK_PER_RUN, test_data_path)
    data = np.ones((len(channel_labels), 1000))
    saver.save_raw_data(data, 500)
    file_path = test_data_path + filename
    saver.edf_data_mark(550, 'OK')
    saver.close_edf_file()
    os.remove(file_path)


def test_edf_data_mark_timestamp_none():
    saver = SigSave(channel_labels, 1000, 375000, -375000)
    subject = create_subject()
    saver.set_edf_header(subject, filename, TASK_PER_RUN, test_data_path)
    data = np.ones((len(channel_labels), 1000))
    saver.save_raw_data(data)
    file_path = test_data_path + filename
    saver.edf_data_mark(0.5, 'OK')
    saver.close_edf_file()
    os.remove(file_path)