# -*- 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 ( ) )