You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
134 lines
6.9 KiB
134 lines
6.9 KiB
# -*- 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())
|