diff --git a/tests/phpunit/MediaWikiParserTest.php b/tests/phpunit/MediaWikiParserTest.php
index 096f337..d10f337 100644
--- a/tests/phpunit/MediaWikiParserTest.php
+++ b/tests/phpunit/MediaWikiParserTest.php
@@ -13,7 +13,7 @@ class MediaWikiParserTest extends MediaWikiTestCase {
 		$title = Title::newFromText( 'test' );
 		$options = new ParserOptions();
 		// Required for MW >= 1.30
-		if( method_exists ( $options, 'setOption' ) ) {
+		if ( method_exists( $options, 'setOption' ) ) {
 			$options->setOption( 'wrapclass', false );
 		}
 		$this->parser->startExternalParse( $title, $options, 'text', true );
diff --git a/tests/phpunit/PortableInfoboxRenderServiceTest.php b/tests/phpunit/PortableInfoboxRenderServiceTest.php
index cc38e3b..638e453 100644
--- a/tests/phpunit/PortableInfoboxRenderServiceTest.php
+++ b/tests/phpunit/PortableInfoboxRenderServiceTest.php
@@ -22,38 +22,18 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 		return $DOM->saveXML();
 	}
 
-	public function testGetImageHelper() {
-		$infoboxRenderService = new PortableInfoboxRenderService();
-
-		$reflection = new ReflectionClass( $infoboxRenderService );
-		$reflection_method = $reflection->getMethod( 'getImageHelper' );
-		$reflection_method->setAccessible( true );
-
-		$this->assertInstanceOf(
-			\PortableInfobox\Helpers\PortableInfoboxImagesHelper::class,
-			$reflection_method->invoke( $infoboxRenderService )
-		);
-	}
-
 	/**
 	 * @covers PortableInfoboxRenderService::renderInfobox
 	 * @covers PortableInfobox\Helpers\PortableInfoboxTemplateEngine
 	 * @param $input
 	 * @param $expectedOutput
 	 * @param $description
-	 * @param $mockParams
 	 * @param $accentColor
 	 * @param $accentColorText
 	 * @dataProvider renderInfoboxDataProvider
 	 */
-	public function testRenderInfobox( $input, $expectedOutput, $description, $mockParams, $accentColor, $accentColorText ) {
+	public function testRenderInfobox( $input, $expectedOutput, $description, $accentColor, $accentColorText ) {
 		$infoboxRenderService = new PortableInfoboxRenderService();
-		$helper = new DummyPIImageHelper( $mockParams );
-
-		$reflection = new ReflectionClass( $infoboxRenderService );
-		$reflection_property = $reflection->getProperty( 'helper' );
-		$reflection_property->setAccessible( true );
-		$reflection_property->setValue( $infoboxRenderService, $helper );
 
 		$actualOutput = $infoboxRenderService->renderInfobox( $input, '', '', $accentColor, $accentColorText );
 		$expectedHtml = $this->normalizeHTML( $expectedOutput );
@@ -68,7 +48,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 				'input' => [],
 				'output' => '',
 				'description' => 'Empty data should yield no infobox markup',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -85,7 +64,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								<h2 class="pi-item pi-item-spacing pi-title">Test Title</h2>
 							</aside>',
 				'description' => 'Only title',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -102,7 +80,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								<h2 class="pi-item pi-item-spacing pi-title" style="background-color:#FFF;color:#000;">Test Title</h2>
 							</aside>',
 				'description' => 'Only title with custom colors',
-				'extendImageData' => [],
 				'accentColor' => '#FFF',
 				'accentColorText' => '#000'
 			],
@@ -114,7 +91,12 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 							[
 								'alt' => 'image alt',
 								'url' => 'http://image.jpg',
-								'caption' => 'Lorem ipsum dolor'
+								'caption' => 'Lorem ipsum dolor',
+								'width' => '400',
+								'height' => '200',
+								'thumbnail' => 'http://thumbnail.jpg',
+								'thumbnail2x' => 'http://thumbnail2x.jpg',
+								'isImage' => true
 							]
 						]
 					]
@@ -129,16 +111,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</figure>
 							</aside>',
 				'description' => 'Only image',
-				'extendImageData' => [
-					'alt' => 'image alt',
-					'url' => 'http://image.jpg',
-					'caption' => 'Lorem ipsum dolor',
-					'width' => '400',
-					'height' => '200',
-					'thumbnail' => 'http://thumbnail.jpg',
-					'thumbnail2x' => 'http://thumbnail2x.jpg',
-					'isImage' => true
-				],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -150,7 +122,8 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 							[
 								'alt' => 'image alt',
 								'url' => 'http://image.jpg',
-								'caption' => 'Lorem ipsum dolor'
+								'caption' => 'Lorem ipsum dolor',
+								'isVideo' => true
 							]
 						]
 					]
@@ -166,12 +139,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</figure>
 							</aside>',
 				'description' => 'Only video',
-				'extendImageData' => [
-					'alt' => 'image alt',
-					'url' => 'http://image.jpg',
-					'caption' => 'Lorem ipsum dolor',
-					'isVideo' => true
-				],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -188,7 +155,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								<nav class="pi-navigation pi-item-spacing pi-secondary-background pi-secondary-font">navigation value</nav>
 							</aside>',
 				'description' => 'navigation only',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -209,7 +175,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</div>
 							</aside>',
 				'description' => 'Only pair',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -226,7 +191,12 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 						'data' => [
 							[
 								'alt' => 'image alt',
-								'url' => 'http://image.jpg'
+								'url' => 'http://image.jpg',
+								'width' => '400',
+								'height' => '200',
+								'thumbnail' => 'http://thumbnail.jpg',
+								'thumbnail2x' => 'http://thumbnail2x.jpg',
+								'isImage' => true
 							]
 						]
 					],
@@ -252,15 +222,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 									</div>
 							</aside>',
 				'description' => 'Simple infobox with title, image and key-value pair',
-				'extendImageData' => [
-					'alt' => 'image alt',
-					'url' => 'http://image.jpg',
-					'width' => '400',
-					'height' => '200',
-					'thumbnail' => 'http://thumbnail.jpg',
-					'thumbnail2x' => 'http://thumbnail2x.jpg',
-					'isImage' => true
-				],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -292,7 +253,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 									</div>
 							</aside>',
 				'description' => 'Simple infobox with title, INVALID image and key-value pair',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -320,7 +280,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</div>
 							</aside>',
 				'description' => 'Simple infobox with title, empty image and key-value pair',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -378,7 +337,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</section>
 							</aside>',
 				'description' => 'Infobox with title, group with header and two key-value pairs',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -436,7 +394,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</section>
 							</aside>',
 				'description' => 'Infobox with title, group with header and two key-value pairs, custom accent color and accent text color',
-				'extendImageData' => [],
 				'accentColor' => '#FFF',
 				'accentColorText' => '#000'
 			],
@@ -498,7 +455,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</section>
 							</aside>',
 				'description' => 'Infobox with horizontal group',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -544,7 +500,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</section>
 							</aside>',
 				'description' => 'Infobox with horizontal group without header and labels',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -563,7 +518,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</nav>
 							</aside>',
 				'description' => 'Infobox with navigation',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -614,7 +568,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</section>
 							</aside>',
 				'description' => 'Horizontal group data without header',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -671,7 +624,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</section>
 							</aside>',
 				'description' => 'Horizontal group data with empty label',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -715,7 +667,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</section>
 							</aside>',
 				'description' => 'Horizontal group data with empty label',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -804,7 +755,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</section>
 							</aside>',
 				'description' => 'Flex wrapped group of 5 elements with row size 3',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -882,7 +832,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</section>
 							</aside>',
 				'description' => 'Flex wrapped group of 4 elements with row size 3',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -921,7 +870,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</section>
 							</aside>',
 				'description' => 'Flex wrapped group of a single element with row size 3',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -971,7 +919,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</section>
 							</aside>',
 				'description' => 'Flex wrapped group of 2 + 1 with row size 3',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -1021,7 +968,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</section>
 							</aside>',
 				'description' => 'Flex wrapped group of 2 + 1 with row size 7',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -1094,7 +1040,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</section>
 							</aside>',
 				'description' => 'Flex wrapped group of 2 + 2 + 2 with row size 3',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -1196,7 +1141,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</section>
 							</aside>',
 				'description' => 'Flex wrapped group of 1 + 1 + default + 1 + 1 + 1 with row size 3',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -1298,7 +1242,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 							  </section>
 							</aside>',
 				'description' => 'Flex wrapped group of 1 + 1 + 1 + 1 + 1 + default with row size 3',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -1365,7 +1308,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 							  </section>
 							</aside>',
 				'description' => 'Flex wrapped group of 40 + 1 + 1 with row size 3',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -1438,7 +1380,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 							  </section>
 							</aside>',
 				'description' => 'Flex wrapped group of 1 + 40 + 1 with row size 3',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -1512,7 +1453,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 							  </section>
 							</aside>',
 				'description' => 'Flex wrapped group of 1 + 1 + title + 1 with row size 3',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -1593,7 +1533,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 							  </section>
 							</aside>',
 				'description' => 'Flex wrapped group of 1 (no label) + 1 (no label) + title + 1 (no label) + 1 with row size 3',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -1611,12 +1550,25 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 							[
 								'alt' => 'image alt',
 								'url' => 'http://image.jpg',
-								'caption' => 'caption'
+								'caption' => 'caption',
+								'ref' => 1,
+								'width' => '400',
+								'height' => '200',
+								'thumbnail' => 'http://thumbnail.jpg',
+								'thumbnail2x' => 'http://thumbnail2x.jpg',
+								'isImage' => true,
+								'isFirst' => true
 							],
 							[
 								'alt' => 'image alt',
 								'url' => 'http://image.jpg',
-								'caption' => 'caption'
+								'caption' => 'caption',
+								'ref' => 2,
+								'width' => '400',
+								'height' => '200',
+								'thumbnail' => 'http://thumbnail.jpg',
+								'thumbnail2x' => 'http://thumbnail2x.jpg',
+								'isImage' => true
 							]
 						]
 					]
@@ -1647,17 +1599,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</div>
 							</aside>',
 				'description' => 'Simple infobox with title and image collection',
-				'extendImageData' => [
-					'alt' => 'image alt',
-					'url' => 'http://image.jpg',
-					'caption' => 'caption',
-					'ref' => 1,
-					'width' => '400',
-					'height' => '200',
-					'thumbnail' => 'http://thumbnail.jpg',
-					'thumbnail2x' => 'http://thumbnail2x.jpg',
-					'isImage' => true
-				],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -1715,7 +1656,6 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</section>
 							</aside>',
 				'description' => 'Infobox with title, collapsible group with header and two key-value pairs',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			],
@@ -1773,32 +1713,9 @@ class PortableInfoboxRenderServiceTest extends MediaWikiTestCase {
 								</section>
 							</aside>',
 				'description' => 'Infobox with title, collapsed group with header and two key-value pairs',
-				'extendImageData' => [],
 				'accentColor' => '',
 				'accentColorText' => ''
 			]
 		];
 	}
 }
-
-class DummyPIImageHelper {
-	private $imageData = [];
-	private $ref = 1;
-
-	public function __construct( $imageData = [] ) {
-		$this->imageData = $imageData;
-	}
-
-	public function extendImageData( $imageData ) {
-		$data = $this->imageData;
-		$data['ref'] = $this->ref++;
-		return $data;
-	}
-
-	public function extendImageCollectionData( $images ) {
-		$images[0]['isFirst'] = true;
-		return [
-			'images' => $images
-		];
-	}
-}
diff --git a/tests/phpunit/helpers/PortableInfoboxImagesHelperTest.php b/tests/phpunit/helpers/PortableInfoboxImagesHelperTest.php
index 5146cd4..3cad086 100644
--- a/tests/phpunit/helpers/PortableInfoboxImagesHelperTest.php
+++ b/tests/phpunit/helpers/PortableInfoboxImagesHelperTest.php
@@ -8,7 +8,6 @@ use PortableInfobox\Helpers\PortableInfoboxImagesHelper;
  */
 class PortableInfoboxImagesHelperTest extends MediaWikiTestCase {
 	private $helper;
-	private static $testCustomWidthLogicCount;
 
 	protected function setUp() {
 		parent::setUp();
@@ -92,17 +91,11 @@ class PortableInfoboxImagesHelperTest extends MediaWikiTestCase {
 	 * @dataProvider customWidthProvider
 	 */
 	public function testCustomWidthLogic( $customWidth, $preferredWidth, $resultDimensions, $thumbnailDimensions, $thumbnail2xDimensions, $originalDimension ) {
-		self::$testCustomWidthLogicCount++;
 		$expected = [
-			'name' => 'test',
-			'ref' => self::$testCustomWidthLogicCount,
 			'thumbnail' => null,
 			'thumbnail2x' => null,
 			'width' => $resultDimensions['width'],
-			'height' => $resultDimensions['height'],
-			'isImage' => true,
-			'isVideo' => false,
-			'isAudio' => false
+			'height' => $resultDimensions['height']
 		];
 		$thumb = $this->getMockBuilder( 'ThumbnailImage' )
 			->disableOriginalConstructor()
@@ -115,22 +108,17 @@ class PortableInfoboxImagesHelperTest extends MediaWikiTestCase {
 		$file->expects( $this->once() )->method( 'exists' )->will( $this->returnValue( true ) );
 		$file->expects( $this->once() )->method( 'getWidth' )->will( $this->returnValue( $originalDimension['width'] ) );
 		$file->expects( $this->once() )->method( 'getHeight' )->will( $this->returnValue( $originalDimension['height'] ) );
-		$file->expects( $this->any() )->method( 'getMediaType' )->will( $this->returnValue( MEDIATYPE_BITMAP ) );
+		$file->expects( $this->once() )->method( 'getMediaType' )->will( $this->returnValue( MEDIATYPE_BITMAP ) );
 
 		$file->expects( $this->any() )
 			->method( 'transform' )
 			->with( $this->logicalOr( $this->equalTo( $thumbnailDimensions ), $this->equalTo( $thumbnail2xDimensions ) ) )
 			->will( $this->returnValue( $thumb ) );
 
-		$helper = $this->getMock( PortableInfoboxImagesHelper::class, [ 'getFileFromTitle' ] );
-		$helper->expects( $this->any() )
-			->method( 'getFileFromTitle' )
-			->will( $this->returnValue( $file ) );
-
 		global $wgPortableInfoboxCustomImageWidth;
 		$wgPortableInfoboxCustomImageWidth = $customWidth;
 
-		$result = $helper->extendImageData( [ 'name' => 'test' ], $preferredWidth );
+		$result = $this->helper->extendImageData( $file, $preferredWidth );
 
 		$this->assertEquals( $expected, $result );
 	}
diff --git a/tests/phpunit/nodes/NodeMediaTest.php b/tests/phpunit/nodes/NodeMediaTest.php
index 0a28290..af388bf 100644
--- a/tests/phpunit/nodes/NodeMediaTest.php
+++ b/tests/phpunit/nodes/NodeMediaTest.php
@@ -1,6 +1,8 @@
 <?php
 
 use PortableInfobox\Helpers\PortableInfoboxDataBag;
+use PortableInfobox\Helpers\PortableInfoboxImagesHelper;
+use PortableInfobox\Parser\Nodes\NodeFactory;
 use PortableInfobox\Parser\Nodes\NodeMedia;
 
 /**
@@ -9,13 +11,6 @@ use PortableInfobox\Parser\Nodes\NodeMedia;
  */
 class NodeMediaTest extends MediaWikiTestCase {
 
-	protected function setUp() {
-		parent::setUp();
-
-		global $wgUseInstantCommons;
-		$wgUseInstantCommons = false;
-	}
-
 	/**
 	 * @covers       PortableInfobox\Parser\Nodes\NodeMedia::getGalleryData
 	 * @covers       PortableInfobox\Helpers\PortableInfoboxDataBag
@@ -139,18 +134,27 @@ class NodeMediaTest extends MediaWikiTestCase {
 	 *
 	 * @param $markup
 	 * @param $params
+	 * @param $mediaType
 	 * @param $expected
 	 */
-	public function testData( $markup, $params, $expected ) {
-		$imageMock = empty( $params ) ? null : new ImageMock();
-		$xmlObj = PortableInfobox\Parser\XmlParser::parseXmlString( $markup );
+	public function testData( $markup, $params, $mediaType, $expected ) {
+		$fileMock = is_null( $mediaType ) ? null : new FileMock( $mediaType );
+		$node = NodeFactory::newFromXML( $markup, $params );
 
-		$mock = $this->getMock( NodeMedia::class, [ 'getFilefromTitle' ], [ $xmlObj, $params ] );
-		$mock->expects( $this->any() )
-			->method( 'getFilefromTitle' )
-			->willReturn( $imageMock );
+		$helper = $this->getMock( PortableInfoboxImagesHelper::class, [ 'getFile', 'extendImageData' ] );
+		$helper->expects( $this->any() )
+			->method( 'getFile' )
+			->willReturn( $fileMock );
+		$helper->expects( $this->any() )
+			->method( 'extendImageData' )
+			->willReturn( [] );
 
-		$this->assertEquals( $expected, $mock->getData() );
+		$reflection = new ReflectionClass( $node );
+		$reflectionProperty = $reflection->getProperty( 'helper' );
+		$reflectionProperty->setAccessible( true );
+		$reflectionProperty->setValue( $node, $helper );
+
+		$this->assertEquals( $expected, $node->getData() );
 	}
 
 	public function dataProvider() {
@@ -159,28 +163,105 @@ class NodeMediaTest extends MediaWikiTestCase {
 			[
 				'<media source="img"></media>',
 				[],
+				null,
 				[ [] ]
 			],
 			[
 				'<media source="img"></media>',
 				[ 'img' => 'test.jpg' ],
-				[ [ 'url' => '', 'name' => 'Test.jpg', 'alt' => 'Test.jpg', 'caption' => null ] ]
+				MEDIATYPE_BITMAP,
+				[ [
+					'url' => 'http://test.url',
+					'name' => 'Test.jpg',
+					'alt' => 'Test.jpg',
+					'caption' => null,
+					'isImage' => true,
+					'isVideo' => false,
+					'isAudio' => false
+				] ]
 			],
 			[
 				'<media source="img"><alt><default>test alt</default></alt></media>',
 				[ 'img' => 'test.jpg' ],
-				[ [ 'url' => '', 'name' => 'Test.jpg', 'alt' => 'test alt', 'caption' => null ] ]
+				MEDIATYPE_BITMAP,
+				[ [
+					'url' => 'http://test.url',
+					'name' => 'Test.jpg',
+					'alt' => 'test alt',
+					'caption' => null,
+					'isImage' => true,
+					'isVideo' => false,
+					'isAudio' => false
+				] ]
 			],
 			[
 				'<media source="img"><alt source="alt source"><default>test alt</default></alt></media>',
 				[ 'img' => 'test.jpg', 'alt source' => 2 ],
-				[ [ 'url' => '', 'name' => 'Test.jpg', 'alt' => 2, 'caption' => null ] ]
+				MEDIATYPE_BITMAP,
+				[ [
+					'url' => 'http://test.url',
+					'name' => 'Test.jpg',
+					'alt' => 2,
+					'caption' => null,
+					'isImage' => true,
+					'isVideo' => false,
+					'isAudio' => false
+				] ]
 			],
 			[
 				'<media source="img"><alt><default>test alt</default></alt><caption source="img"/></media>',
 				[ 'img' => 'test.jpg' ],
-				[ [ 'url' => '', 'name' => 'Test.jpg', 'alt' => 'test alt', 'caption' => 'test.jpg' ] ]
+				MEDIATYPE_BITMAP,
+				[ [
+					'url' => 'http://test.url',
+					'name' => 'Test.jpg',
+					'alt' => 'test alt',
+					'caption' => 'test.jpg',
+					'isImage' => true,
+					'isVideo' => false,
+					'isAudio' => false
+				] ]
 			],
+			[
+				'<media source="media" />',
+				[ 'media' => 'test.webm' ],
+				MEDIATYPE_VIDEO,
+				[ [
+					'url' => 'http://test.url',
+					'name' => 'Test.webm',
+					'alt' => 'Test.webm',
+					'caption' => null,
+					'isImage' => false,
+					'isVideo' => true,
+					'isAudio' => false
+				] ]
+			],
+			[
+				'<media source="media" video="false" />',
+				[ 'media' => 'test.webm' ],
+				MEDIATYPE_VIDEO,
+				[ [] ]
+			],
+			[
+				'<media source="media" />',
+				[ 'media' => 'test.ogg' ],
+				MEDIATYPE_AUDIO,
+				[ [
+					'url' => 'http://test.url',
+					'name' => 'Test.ogg',
+					'alt' => 'Test.ogg',
+					'caption' => null,
+					'isImage' => false,
+					'isVideo' => false,
+					'isAudio' => true
+				] ]
+			],
+			[
+				'<media source="media" audio="false" />',
+				[ 'media' => 'test.ogg' ],
+				MEDIATYPE_AUDIO,
+				[ [] ]
+			]
 		];
 	}
 
@@ -193,7 +274,7 @@ class NodeMediaTest extends MediaWikiTestCase {
 	 * @param $expected
 	 */
 	public function testIsEmpty( $markup, $params, $expected ) {
-		$node = PortableInfobox\Parser\Nodes\NodeFactory::newFromXML( $markup, $params );
+		$node = NodeFactory::newFromXML( $markup, $params );
 
 		$this->assertEquals( $expected, $node->isEmpty() );
 	}
@@ -212,7 +293,7 @@ class NodeMediaTest extends MediaWikiTestCase {
 	 * @param $expected
 	 */
 	public function testSources( $markup, $expected ) {
-		$node = PortableInfobox\Parser\Nodes\NodeFactory::newFromXML( $markup, [] );
+		$node = NodeFactory::newFromXML( $markup, [] );
 
 		$this->assertEquals( $expected, $node->getSources() );
 	}
@@ -247,7 +328,7 @@ class NodeMediaTest extends MediaWikiTestCase {
 
 	/** @dataProvider metadataProvider */
 	public function testMetadata( $markup, $expected ) {
-		$node = PortableInfobox\Parser\Nodes\NodeFactory::newFromXML( $markup, [] );
+		$node = NodeFactory::newFromXML( $markup, [] );
 
 		$this->assertEquals( $expected, $node->getMetadata() );
 	}
@@ -265,88 +346,6 @@ class NodeMediaTest extends MediaWikiTestCase {
 		];
 	}
 
-	/**
-	 * @dataProvider videoProvider
-	 * @param $markup
-	 * @param $params
-	 * @param $expected
-	 * @throws PortableInfobox\Parser\XmlMarkupParseErrorException
-	 */
-	public function testVideo( $markup, $params, $expected ) {
-		$videoMock = new VideoMock();
-		$xmlObj = PortableInfobox\Parser\XmlParser::parseXmlString( $markup );
-
-		$mock = $this->getMock( NodeMedia::class, [ 'getFilefromTitle' ], [ $xmlObj, $params ] );
-		$mock->expects( $this->any() )
-			->method( 'getFilefromTitle' )
-			->willReturn( $videoMock );
-
-		$this->assertEquals( $expected, $mock->getData() );
-	}
-
-	public function videoProvider() {
-		return [
-			[
-				'<media source="media" />',
-				[ 'media' => 'test.webm' ],
-				[
-					[
-						'url' => 'http://test.url',
-						'name' => 'Test.webm',
-						'alt' => 'Test.webm',
-						'caption' => null
-					]
-				]
-			],
-			[
-				'<media source="media" video="false" />',
-				[ 'media' => 'test.webm' ],
-				[ [] ]
-			]
-		];
-	}
-
-	/**
-	 * @dataProvider audioProvider
-	 * @param $markup
-	 * @param $params
-	 * @param $expected
-	 * @throws PortableInfobox\Parser\XmlMarkupParseErrorException
-	 */
-	public function testAudio( $markup, $params, $expected ) {
-		$audioMock = new AudioMock();
-		$xmlObj = PortableInfobox\Parser\XmlParser::parseXmlString( $markup );
-
-		$mock = $this->getMock( NodeMedia::class, [ 'getFilefromTitle' ], [ $xmlObj, $params ] );
-		$mock->expects( $this->any() )
-			->method( 'getFilefromTitle' )
-			->willReturn( $audioMock );
-
-		$this->assertEquals( $expected, $mock->getData() );
-	}
-
-	public function audioProvider() {
-		return [
-			[
-				'<media source="media" />',
-				[ 'media' => 'test.ogg' ],
-				[
-					[
-						'url' => 'http://test.url',
-						'name' => 'Test.ogg',
-						'alt' => 'Test.ogg',
-						'caption' => null
-					]
-				]
-			],
-			[
-				'<media source="media" audio="false" />',
-				[ 'media' => 'test.ogg' ],
-				[ [] ]
-			]
-		];
-	}
-
 	/**
 	 * @covers PortableInfobox\Parser\Nodes\NodeMedia::isTypeAllowed
 	 * @covers PortableInfobox\Parser\Nodes\NodeAudio
@@ -413,29 +412,15 @@ class NodeMediaTest extends MediaWikiTestCase {
 	}
 }
 
-class ImageMock {
+class FileMock {
+	protected $mediaType;
+
+	public function __construct( $mediaType ) {
+		$this->mediaType = $mediaType;
+	}
+
 	public function getMediaType() {
-		return MEDIATYPE_BITMAP;
-	}
-
-	public function getUrl() {
-		return '';
-	}
-}
-
-class VideoMock {
-	public function getMediaType() {
-		return MEDIATYPE_VIDEO;
-	}
-
-	public function getUrl() {
-		return 'http://test.url';
-	}
-}
-
-class AudioMock {
-	public function getMediaType() {
-		return MEDIATYPE_AUDIO;
+		return $this->mediaType;
 	}
 
 	public function getUrl() {