bazarr/libs/enzyme/tests/test_parsers.py

# -*- coding: utf-8 -*-
from enzyme.parsers import ebml
import io
import os.path
import requests
import unittest
import yaml
import zipfile


# Test directory
TEST_DIR = os.path.join(os.path.dirname(__file__), os.path.splitext(__file__)[0])

# EBML validation directory
EBML_VALIDATION_DIR = os.path.join(os.path.dirname(__file__), 'parsers', 'ebml')


def setUpModule():
    if not os.path.exists(TEST_DIR):
        r = requests.get('http://downloads.sourceforge.net/project/matroska/test_files/matroska_test_w1_1.zip')
        with zipfile.ZipFile(io.BytesIO(r.content), 'r') as f:
            f.extractall(TEST_DIR)


class EBMLTestCase(unittest.TestCase):
    def setUp(self):
        self.stream = io.open(os.path.join(TEST_DIR, 'test1.mkv'), 'rb')
        with io.open(os.path.join(EBML_VALIDATION_DIR, 'test1.mkv.yml'), 'r') as yml:
            self.validation = yaml.safe_load(yml)
        self.specs = ebml.get_matroska_specs()

    def tearDown(self):
        self.stream.close()

    def check_element(self, element_id, element_type, element_name, element_level, element_position, element_size, element_data, element,
                      ignore_element_types=None, ignore_element_names=None, max_level=None, include_element_names=None):
        """Recursively check an element"""
        # base
        self.assertTrue(element.id == element_id)
        self.assertTrue(element.type == element_type)
        self.assertTrue(element.name == element_name)
        self.assertTrue(element.level == element_level)
        self.assertTrue(element.position == element_position)
        self.assertTrue(element.size == element_size)
        # Element
        if not isinstance(element_data, list):
            self.assertTrue(type(element) == ebml.Element)
            if element_type != ebml.BINARY:
                self.assertTrue(element.data == element_data)
            return
        # MasterElement
        if ignore_element_types is not None:  # filter validation on element types
            element_data = [e for e in element_data if e[1] not in ignore_element_types]
        if ignore_element_names is not None:  # filter validation on element names
            element_data = [e for e in element_data if e[2] not in ignore_element_names]
        if include_element_names is not None:  # filter validation on element names
            element_data = [e for e in element_data if e[2] in include_element_names]
        if element.level == max_level:  # special check when maximum level is reached
            self.assertTrue(element.data is None)
            return
        self.assertTrue(len(element.data) == len(element_data))
        for i in range(len(element.data)):
            self.check_element(element_data[i][0], element_data[i][1], element_data[i][2], element_data[i][3],
                               element_data[i][4], element_data[i][5], element_data[i][6], element.data[i], ignore_element_types,
                               ignore_element_names, max_level,include_element_names)

    def test_parse_full(self):
        result = ebml.parse(self.stream, self.specs)
        self.assertTrue(len(result) == len(self.validation))
        for i in range(len(self.validation)):
            self.check_element(self.validation[i][0], self.validation[i][1], self.validation[i][2], self.validation[i][3],
                               self.validation[i][4], self.validation[i][5], self.validation[i][6], result[i])

    def test_parse_ignore_element_types(self):
        ignore_element_types = [ebml.INTEGER, ebml.BINARY]
        result = ebml.parse(self.stream, self.specs, ignore_element_types=ignore_element_types)
        self.validation = [e for e in self.validation if e[1] not in ignore_element_types]
        self.assertTrue(len(result) == len(self.validation))
        for i in range(len(self.validation)):
            self.check_element(self.validation[i][0], self.validation[i][1], self.validation[i][2], self.validation[i][3],
                               self.validation[i][4], self.validation[i][5], self.validation[i][6], result[i], ignore_element_types=ignore_element_types)

    def test_parse_ignore_element_names(self):
        ignore_element_names = ['EBML', 'SimpleBlock']
        result = ebml.parse(self.stream, self.specs, ignore_element_names=ignore_element_names)
        self.validation = [e for e in self.validation if e[2] not in ignore_element_names]
        self.assertTrue(len(result) == len(self.validation))
        for i in range(len(self.validation)):
            self.check_element(self.validation[i][0], self.validation[i][1], self.validation[i][2], self.validation[i][3],
                               self.validation[i][4], self.validation[i][5], self.validation[i][6], result[i], ignore_element_names=ignore_element_names)

    def test_parse_include_element_names(self):
        include_element_names = ['Segment','Cluster']
        result = ebml.parse(self.stream, self.specs, include_element_names=include_element_names)
        self.validation = [e for e in self.validation if e[2] in include_element_names]
        self.assertTrue(len(result) == len(self.validation))
        for i in range(len(self.validation)):
            self.check_element(self.validation[i][0], self.validation[i][1], self.validation[i][2], self.validation[i][3],
                               self.validation[i][4], self.validation[i][5], self.validation[i][6], result[i], include_element_names=include_element_names)
            
    def test_parse_max_level(self):
        max_level = 3
        result = ebml.parse(self.stream, self.specs, max_level=max_level)
        self.validation = [e for e in self.validation if e[3] <= max_level]
        self.assertTrue(len(result) == len(self.validation))
        for i in range(len(self.validation)):
            self.check_element(self.validation[i][0], self.validation[i][1], self.validation[i][2], self.validation[i][3],
                               self.validation[i][4], self.validation[i][5], self.validation[i][6], result[i], max_level=max_level)


def generate_yml(filename, specs):
    """Generate  a validation file for the test video"""
    def _to_builtin(elements):
        """Recursively convert elements to built-in types"""
        result = []
        for e in elements:
            if isinstance(e, ebml.MasterElement):
                result.append((e.id, e.type, e.name, e.level, e.position, e.size, _to_builtin(e.data)))
            else:
                result.append((e.id, e.type, e.name, e.level, e.position, e.size, None if isinstance(e.data, io.BytesIO) else e.data))
        return result
    video = io.open(os.path.join(TEST_DIR, filename), 'rb')
    yml = io.open(os.path.join(EBML_VALIDATION_DIR, filename + '.yml'), 'w')
    yaml.safe_dump(_to_builtin(ebml.parse(video, specs)), yml)


def suite():
    suite = unittest.TestSuite()
    suite.addTest(unittest.TestLoader().loadTestsFromTestCase(EBMLTestCase))
    return suite

if __name__ == '__main__':
    unittest.TextTestRunner().run(suite())