My first crack at a ConnectionParser. Tests pass but aren't pretty yet. I'm not sure if I like this approach or not.

t8y8 · t8y8 · commit e0ceed56a4ff · 2016-05-12T21:54:48.000-07:00
diff --git a/tableaudocumentapi/datasource.py b/tableaudocumentapi/datasource.py
@@ -6,6 +6,26 @@
 import xml.etree.ElementTree as ET
 from tableaudocumentapi import Connection
 
+class ConnectionParser(object):
+
+    def __init__(self, datasource_xml, version):
+        self._dsxml = datasource_xml
+        self._dsversion = version
+
+    def get_connections(self):
+        if float(self._dsversion) < 10:
+            connections = self._extract_legacy_connection()
+        else:
+            connections = self._extract_federated_connections()
+        return connections
+
+    def _extract_federated_connections(self):
+        return list(map(Connection,self._dsxml.findall('.//named-connections/named-connection/*')))
+
+    def _extract_legacy_connection(self):
+        return Connection(self._dsxml.find('connection'))
+
+
 class Datasource(object):
     """
     A class for writing datasources to Tableau files.
@@ -27,10 +47,8 @@ def __init__(self, dsxml, filename=None):
         self._datasourceTree = ET.ElementTree(self._datasourceXML)
         self._name = self._datasourceXML.get('name') or self._datasourceXML.get('formatted-name') # TDS files don't have a name attribute
         self._version = self._datasourceXML.get('version')
-        if self._version == '10.0':
-            self._connection = list(map(Connection,self._datasourceXML.findall('.//named-connections/named-connection/*')))
-        else:
-            self._connection = Connection(self._datasourceXML.find('connection'))
+        self._connection_parser = ConnectionParser(self._datasourceXML, version=self._version)
+        self._connection = self._connection_parser.get_connections()
 
     @classmethod
     def from_file(cls, filename):
diff --git a/test.py b/test.py
@@ -39,6 +39,13 @@ def test_is_valid_file_with_invalid_inputs(self):
         self.assertFalse(Workbook._is_valid_file('file1.tds2'))
         self.assertFalse(Workbook._is_valid_file('file2.twb3'))
 
+class ConnectionParserTests(unittest.TestCase):
+
+    def test_can_extract_legacy_connection(self):
+        pass
+
+    def test_can_extract_federated_connections(self):
+        pass
 
 class ConnectionModelTests(unittest.TestCase):
 
@@ -128,6 +135,7 @@ def test_can_update_datasource_connection_and_saveV10(self):
         self.assertEqual(new_wb.datasources[0].connection[0].dbname, 'newdb.test.tsi.lan')
 
         temp.close()
+        os.unlink(temp.name)
 
     def test_can_extract_datasourceV10(self):
         temp = io.FileIO('v10test.twb', 'w')
