From c318180f48376b4b5afca1668ab71749658fa2f0 Mon Sep 17 00:00:00 2001
From: zhangshine <zhangshine0125@gmail.com>
Date: Thu, 11 Jun 2015 21:23:28 +0800
Subject: [PATCH] test Service provider

---
 pyinvoice/models.py     |  2 +-
 pyinvoice/templates.py  | 21 ++++++++++-----------
 tests/test_templates.py | 40 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 51 insertions(+), 12 deletions(-)

diff --git a/pyinvoice/models.py b/pyinvoice/models.py
index 0ed6630..0c93496 100644
--- a/pyinvoice/models.py
+++ b/pyinvoice/models.py
@@ -72,7 +72,7 @@ class ServiceProviderInfo(AddressInfo):
         :type state: str or unicode or None
         :type country: str or unicode or None
         :type post_code: str or unicode or None
-        :type vat_tax_number: str or unicode or None
+        :type vat_tax_number: str or unicode or int or None
         """
         super(ServiceProviderInfo, self).__init__(name, street, city, state, country, post_code)
         self.vat_tax_number = vat_tax_number
diff --git a/pyinvoice/templates.py b/pyinvoice/templates.py
index 0fe096c..eedc1ce 100644
--- a/pyinvoice/templates.py
+++ b/pyinvoice/templates.py
@@ -109,21 +109,20 @@ class SimpleInvoice(SimpleDocTemplate):
             self._story.append(SimpleTable(invoice_info_data, horizontal_align='RIGHT'))
 
     def _service_provider_data(self):
-        props = [('name', 'Name'), ('street', 'Street'), ('city', 'City'), ('state', 'State'),
-                 ('country', 'Country'), ('post_code', 'Post code'), ('vat_tax_number', 'Vat/Tax number')]
+        if isinstance(self.service_provider_info, ServiceProviderInfo):
+            props = [('name', 'Name'), ('street', 'Street'), ('city', 'City'), ('state', 'State'),
+                     ('country', 'Country'), ('post_code', 'Post code'), ('vat_tax_number', 'Vat/Tax number')]
 
-        return self._attribute_to_table_data(self.service_provider_info, props)
+            return self._attribute_to_table_data(self.service_provider_info, props)
+
+        return []
 
     def _build_service_provider_info(self):
         # Merchant
-        if isinstance(self.service_provider_info, ServiceProviderInfo):
-            self._story.append(
-                Paragraph('Merchant', self._defined_styles.get('RightHeading1'))
-            )
-
-            self._story.append(
-                SimpleTable(self._service_provider_data(), horizontal_align='RIGHT')
-            )
+        service_provider_info_data = self._service_provider_data()
+        if service_provider_info_data:
+            self._story.append(Paragraph('Merchant', self._defined_styles.get('RightHeading1')))
+            self._story.append(SimpleTable(service_provider_info_data, horizontal_align='RIGHT'))
 
     def _client_info_data(self):
         props = [('name', 'Name'), ('street', 'Street'), ('city', 'City'), ('state', 'State'),
diff --git a/tests/test_templates.py b/tests/test_templates.py
index 9414f0d..6d8807a 100644
--- a/tests/test_templates.py
+++ b/tests/test_templates.py
@@ -170,4 +170,44 @@ class TestSimpleInvoice(unittest.TestCase):
 
         invoice.finish()
 
+        self.assertTrue(os.path.exists(invoice_path))
+
+    def test_service_provider_info(self):
+        invoice_path = os.path.join(self.file_base_dir, 'service_provider_info.pdf')
+        if os.path.exists(invoice_path):
+            os.remove(invoice_path)
+
+        invoice = SimpleInvoice(invoice_path)
+
+        # Before add service provider info
+        info_data = invoice._service_provider_data()
+        self.assertEqual(info_data, [])
+
+        # Empty info
+        invoice.service_provider_info = ServiceProviderInfo()
+        info_data = invoice._service_provider_data()
+        self.assertEqual(info_data, [])
+
+        invoice.service_provider_info = ServiceProviderInfo(
+            name='CiCiApp',
+            street='Street xxx',
+            city='City ccc',
+            state='State sss',
+            country='Country rrr',
+            post_code='Post code ppp',
+            vat_tax_number=666
+        )
+
+        # After add service provider info
+        info_data = invoice._service_provider_data()
+        self.assertEqual(len(info_data), 7)
+        self.assertEqual(info_data[0][0], 'Name:')
+        self.assertEqual(info_data[0][1], 'CiCiApp')
+        self.assertEqual(info_data[4][0], 'Country:')
+        self.assertEqual(info_data[4][1], 'Country rrr')
+        self.assertEqual(info_data[6][0], 'Vat/Tax number:')
+        self.assertEqual(info_data[6][1], 666)
+
+        invoice.finish()
+
         self.assertTrue(os.path.exists(invoice_path))
\ No newline at end of file