From 6dccb3927aaaf632590c9b80507944798cb1e67e Mon Sep 17 00:00:00 2001
From: zhangshine <zhangshine0125@gmail.com>
Date: Thu, 11 Jun 2015 21:52:06 +0800
Subject: [PATCH] test client info

---
 pyinvoice/models.py     |  2 +-
 pyinvoice/templates.py  | 28 +++++++++++++++++-----------
 tests/test_templates.py | 41 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 12 deletions(-)

diff --git a/pyinvoice/models.py b/pyinvoice/models.py
index 0c93496..49bbf5e 100644
--- a/pyinvoice/models.py
+++ b/pyinvoice/models.py
@@ -92,7 +92,7 @@ class ClientInfo(AddressInfo):
         :type country: str or unicode or None
         :type post_code: str or unicode or None
         :type email: str or unicode or None
-        :type post_code: str or unicode or None
+        :type client_id: str or unicode or int or None
         """
         super(ClientInfo, self).__init__(name, street, city, state, country, post_code)
         self.email = email
diff --git a/pyinvoice/templates.py b/pyinvoice/templates.py
index eedc1ce..99df5ad 100644
--- a/pyinvoice/templates.py
+++ b/pyinvoice/templates.py
@@ -8,7 +8,6 @@ from reportlab.lib.pagesizes import letter
 from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
 from reportlab.lib.units import inch
 from reportlab.platypus import SimpleDocTemplate, Paragraph, Table, Spacer
-from reportlab.platypus.doctemplate import _doNothing
 
 from pyinvoice.components import SimpleTable, TableWithHeader, PaidStamp
 from pyinvoice.models import PDFInfo, Item, Transaction, InvoiceInfo, ServiceProviderInfo, ClientInfo
@@ -125,26 +124,29 @@ class SimpleInvoice(SimpleDocTemplate):
             self._story.append(SimpleTable(service_provider_info_data, horizontal_align='RIGHT'))
 
     def _client_info_data(self):
+        if not isinstance(self.client_info, ClientInfo):
+            return []
+
         props = [('name', 'Name'), ('street', 'Street'), ('city', 'City'), ('state', 'State'),
                  ('country', 'Country'), ('post_code', 'Post code'), ('email', 'Email'), ('client_id', 'Client id')]
         return self._attribute_to_table_data(self.client_info, props)
 
     def _build_client_info(self):
         # ClientInfo
-        if isinstance(self.client_info, ClientInfo):
-            self._story.append(
-                Paragraph('Client', self._defined_styles.get('Heading1'))
-            )
-
-            self._story.append(
-                SimpleTable(self._client_info_data(), horizontal_align='LEFT')
-            )
+        client_info_data = self._client_info_data()
+        if client_info_data:
+            self._story.append(Paragraph('Client', self._defined_styles.get('Heading1')))
+            self._story.append(SimpleTable(client_info_data, horizontal_align='LEFT'))
 
     def _build_service_provider_and_client_info(self):
         if isinstance(self.service_provider_info, ServiceProviderInfo) and isinstance(self.client_info, ClientInfo):
             # Merge Table
             table_data = [
-                [Paragraph('Service Provider', self._defined_styles.get('Heading1')), '', '', Paragraph('Client', self._defined_styles.get('Heading1')), '']
+                [
+                    Paragraph('Service Provider', self._defined_styles.get('Heading1')), '',
+                    '',
+                    Paragraph('Client', self._defined_styles.get('Heading1')), ''
+                ]
             ]
             table_style = [
                 ('SPAN', (0, 0), (1, 0)),
@@ -288,4 +290,8 @@ class SimpleInvoice(SimpleDocTemplate):
         self._build_transactions()
         self._build_bottom_tip()
 
-        self.build(self._story, onFirstPage=PaidStamp(7 * inch, 5.8 * inch) if self.is_paid else _doNothing)
\ No newline at end of file
+        kwargs = {}
+        if self.is_paid:
+            kwargs['onFirstPage'] = PaidStamp(7 * inch, 5.8 * inch)
+
+        self.build(self._story, **kwargs)
\ No newline at end of file
diff --git a/tests/test_templates.py b/tests/test_templates.py
index 6d8807a..e355e95 100644
--- a/tests/test_templates.py
+++ b/tests/test_templates.py
@@ -210,4 +210,45 @@ class TestSimpleInvoice(unittest.TestCase):
 
         invoice.finish()
 
+        self.assertTrue(os.path.exists(invoice_path))
+
+    def test_client_info(self):
+        invoice_path = os.path.join(self.file_base_dir, 'client_info.pdf')
+        if os.path.exists(invoice_path):
+            os.remove(invoice_path)
+
+        invoice = SimpleInvoice(invoice_path)
+
+        # Before add client info
+        info_data = invoice._client_info_data()
+        self.assertEqual(info_data, [])
+
+        # Empty info
+        invoice.client_info = ClientInfo()
+        info_data = invoice._client_info_data()
+        self.assertEqual(info_data, [])
+
+        invoice.client_info = ClientInfo(
+            name='Client ccc',
+            street='Street sss',
+            city='City ccc',
+            state='State sss',
+            country='Country ccc',
+            post_code='Post code ppp',
+            email='Email@example.com',
+            client_id=3214
+        )
+
+        # After add client info
+        info_data = invoice._client_info_data()
+        self.assertEqual(len(info_data), 8)
+        self.assertEqual(info_data[0][0], 'Name:')
+        self.assertEqual(info_data[0][1], 'Client ccc')
+        self.assertEqual(info_data[6][0], 'Email:')
+        self.assertEqual(info_data[6][1], 'Email@example.com')
+        self.assertEqual(info_data[7][0], 'Client id:')
+        self.assertEqual(info_data[7][1], 3214)
+
+        invoice.finish()
+
         self.assertTrue(os.path.exists(invoice_path))
\ No newline at end of file