From 53ca9305deccd7db0c49cd0c152e29df229b8d2e Mon Sep 17 00:00:00 2001
From: zhangshine <zhangshine0125@gmail.com>
Date: Thu, 11 Jun 2015 21:06:35 +0800
Subject: [PATCH] test invoice info

---
 pyinvoice/templates.py  | 18 ++++++++++--------
 tests/test_templates.py | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/pyinvoice/templates.py b/pyinvoice/templates.py
index 5a148b6..0fe096c 100644
--- a/pyinvoice/templates.py
+++ b/pyinvoice/templates.py
@@ -93,18 +93,20 @@ class SimpleInvoice(SimpleDocTemplate):
 
         return data
 
-    def _build_invoice_info(self):
+    def _invoice_info_data(self):
         if isinstance(self.invoice_info, InvoiceInfo):
-            self._story.append(
-                Paragraph('Invoice', self._defined_styles.get('RightHeading1'))
-            )
-
             props = [('invoice_id', 'Invoice id'), ('invoice_datetime', 'Invoice date'),
                      ('due_datetime', 'Invoice due date')]
 
-            self._story.append(
-                SimpleTable(self._attribute_to_table_data(self.invoice_info, props), horizontal_align='RIGHT')
-            )
+            return self._attribute_to_table_data(self.invoice_info, props)
+
+        return []
+
+    def _build_invoice_info(self):
+        invoice_info_data = self._invoice_info_data()
+        if invoice_info_data:
+            self._story.append(Paragraph('Invoice', self._defined_styles.get('RightHeading1')))
+            self._story.append(SimpleTable(invoice_info_data, horizontal_align='RIGHT'))
 
     def _service_provider_data(self):
         props = [('name', 'Name'), ('street', 'Street'), ('city', 'City'), ('state', 'State'),
diff --git a/tests/test_templates.py b/tests/test_templates.py
index d381c76..9414f0d 100644
--- a/tests/test_templates.py
+++ b/tests/test_templates.py
@@ -99,7 +99,7 @@ class TestSimpleInvoice(unittest.TestCase):
         self.assertTrue(os.path.exists(invoice_path))
 
     def test_only_items_with_tax_rate(self):
-        invoice_path = os.path.join(self.file_base_dir, 'only_items.pdf')
+        invoice_path = os.path.join(self.file_base_dir, 'only_items_with_tax.pdf')
         if os.path.exists(invoice_path):
             os.remove(invoice_path)
 
@@ -141,4 +141,33 @@ class TestSimpleInvoice(unittest.TestCase):
 
         invoice.finish()
 
+        self.assertTrue(os.path.exists(invoice_path))
+
+    def test_invoice_info(self):
+        invoice_path = os.path.join(self.file_base_dir, 'invoice_info.pdf')
+        if os.path.exists(invoice_path):
+            os.remove(invoice_path)
+
+        invoice = SimpleInvoice(invoice_path)
+
+        # Before add invoice info
+        invoice_info_data = invoice._invoice_info_data()
+        self.assertEqual(invoice_info_data, [])
+
+        invoice.invoice_info = InvoiceInfo(12)
+
+        # After add invoice info
+        invoice_info_data = invoice._invoice_info_data()
+        self.assertEqual(len(invoice_info_data), 1)
+        self.assertEqual(invoice_info_data[0][0], 'Invoice id:')
+        self.assertEqual(invoice_info_data[0][1], 12)
+
+        invoice.invoice_info = InvoiceInfo(12, invoice_datetime=datetime(2015, 6, 1))
+        invoice_info_data = invoice._invoice_info_data()
+        self.assertEqual(len(invoice_info_data), 2)
+        self.assertEqual(invoice_info_data[1][0], 'Invoice date:')
+        self.assertEqual(invoice_info_data[1][1], '2015-06-01 00:00')
+
+        invoice.finish()
+
         self.assertTrue(os.path.exists(invoice_path))
\ No newline at end of file