Make sure we are loading the expected filetype (#62)

* Proper validation that Workbooks load 'twb(x)' and Datasources load 'tds(x)' files only
* Small cleanup, move the try inside the with

* Add some tests to make sure I don't break anything and cleanup a few small bits

Author: t8y8

tableaudocumentapi/datasource.py

@@ -113,7 +113,7 @@ def __init__(self, dsxml, filename=None):
     def from_file(cls, filename):
         """Initialize datasource from file (.tds)"""
 
-        dsxml = xml_open(filename).getroot()
+        dsxml = xml_open(filename, cls.__name__.lower()).getroot()
         return cls(dsxml, filename)
 
     def save(self):

tableaudocumentapi/workbook.py

@@ -32,7 +32,7 @@ def __init__(self, filename):
 
         self._filename = filename
 
-        self._workbookTree = xml_open(self._filename)
+        self._workbookTree = xml_open(self._filename, self.__class__.__name__.lower())
 
         self._workbookRoot = self._workbookTree.getroot()
         # prepare our datasource objects

tableaudocumentapi/xfile.py

@@ -17,15 +17,28 @@ class TableauVersionNotSupportedException(Exception):
     pass
 
 
-def xml_open(filename):
-    # Determine if this is a twb or twbx and get the xml root
+class TableauInvalidFileException(Exception):
+    pass
+
+
+def xml_open(filename, expected_root=None):
+
     if zipfile.is_zipfile(filename):
         tree = get_xml_from_archive(filename)
     else:
         tree = ET.parse(filename)
-    file_version = Version(tree.getroot().attrib.get('version', '0.0'))
+
+    tree_root = tree.getroot()
+
+    file_version = Version(tree_root.attrib.get('version', '0.0'))
+
     if file_version < MIN_SUPPORTED_VERSION:
         raise TableauVersionNotSupportedException(file_version)
+
+    if expected_root and (expected_root != tree_root.tag):
+        raise TableauInvalidFileException(
+            "'{}'' is not a valid '{}' file".format(filename, expected_root))
+
     return tree
 
 
@@ -40,14 +53,13 @@ def temporary_directory(*args, **kwargs):
 
 def find_file_in_zip(zip_file):
     for filename in zip_file.namelist():
-        try:
-            with zip_file.open(filename) as xml_candidate:
-                ET.parse(xml_candidate).getroot().tag in (
-                    'workbook', 'datasource')
+        with zip_file.open(filename) as xml_candidate:
+            try:
+                ET.parse(xml_candidate)
                 return filename
-        except ET.ParseError:
-            # That's not an XML file by gosh
-            pass
+            except ET.ParseError:
+                # That's not an XML file by gosh
+                pass
 
 
 def get_xml_from_archive(filename):

test/assets/TABLEAU_82_TWB.twb

@@ -0,0 +1 @@
+<?xml version='1.0' encoding='utf-8' ?><workbook source-build='9.3.1 (9300.16.0510.0100)' source-platform='mac' version='8.2' xmlns:user='http://www.tableausoftware.com/xml/user'><datasources><datasource caption='xy (TestV1)' inline='true' name='sqlserver.17u3bqc16tjtxn14e2hxh19tyvpo' version='9.3'><connection authentication='sspi' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username=''></connection></datasource></datasources></workbook>

test/bvt.py

@@ -4,9 +4,12 @@
 import xml.etree.ElementTree as ET
 
 from tableaudocumentapi import Workbook, Datasource, Connection, ConnectionParser
+from tableaudocumentapi.xfile import TableauInvalidFileException, TableauVersionNotSupportedException
 
 TEST_DIR = os.path.dirname(__file__)
 
+TABLEAU_82_TWB = os.path.join(TEST_DIR, 'assets', 'TABLEAU_82_TWB.twb')
+
 TABLEAU_93_TWB = os.path.join(TEST_DIR, 'assets', 'TABLEAU_93_TWB.twb')
 
 TABLEAU_93_TDS = os.path.join(TEST_DIR, 'assets', 'TABLEAU_93_TDS.tds')
@@ -289,10 +292,28 @@ def test_can_open_twbx_and_save_as_changes(self):
 
 
 class EmptyWorkbookWillLoad(unittest.TestCase):
+
     def test_no_exceptions_thrown(self):
         wb = Workbook(EMPTY_WORKBOOK)
         self.assertIsNotNone(wb)
 
 
+class LoadOnlyValidFileTypes(unittest.TestCase):
+
+    def test_exception_when_workbook_given_tdsx(self):
+        with self.assertRaises(TableauInvalidFileException):
+            wb = Workbook(TABLEAU_10_TDSX)
+
+    def test_exception_when_datasource_given_twbx(self):
+        with self.assertRaises(TableauInvalidFileException):
+            ds = Datasource.from_file(TABLEAU_10_TWBX)
+
+
+class SupportedWorkbookVersions(unittest.TestCase):
+
+    def test_82_workbook_throws_exception(self):
+        with self.assertRaises(TableauVersionNotSupportedException):
+            wb = Workbook(TABLEAU_82_TWB)
+
 if __name__ == '__main__':
     unittest.main()