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