server-tools/sql_export_excel/tests/test_sql_query_excel.py

# Copyright (C) 2019 Akretion (<http://www.akretion.com>)
# @author: Florian da Costa
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl.html).

from odoo.tests.common import TransactionCase
import base64
from io import BytesIO
import logging

_logger = logging.getLogger(__name__)

try:
    import openpyxl
except ImportError:
    _logger.debug('Can not import openpyxl')


class TestExportSqlQueryExcel(TransactionCase):

    def setUp(self):
        super().setUp()
        self.wizard_obj = self.env['sql.file.wizard']

    def get_workbook_from_query(self, wizard):
        wizard.export_sql()
        decoded_data = base64.b64decode(wizard.binary_file)
        xlsx_file = BytesIO(decoded_data)
        return openpyxl.load_workbook(xlsx_file)

    def test_excel_file_generation(self):
        test_query = "SELECT 'testcol1' as firstcol, 2 as second_col"
        query_vals = {
            'name': 'Test Query Excel',
            'query': test_query,
            'file_format': 'excel'
        }
        query = self.env['sql.export'].create(query_vals)
        query.button_validate_sql_expression()
        wizard = self.wizard_obj.create({
            'sql_export_id': query.id,
        })
        workbook = self.get_workbook_from_query(wizard)
        ws = workbook.active
        # Check values, header should be here by default
        self.assertEqual(ws.cell(row=1, column=1).value, 'firstcol')
        self.assertEqual(ws.cell(row=2, column=1).value, 'testcol1')
        self.assertEqual(ws.cell(row=2, column=2).value, 2)

        query.write({'header': False})
        wb2 = self.get_workbook_from_query(wizard)
        ws2 = wb2.active
        # Check values, the header should not be present
        self.assertEqual(ws2.cell(row=1, column=1).value, 'testcol1')
        self.assertEqual(ws2.cell(row=1, column=2).value, 2)

    def test_excel_file_insert(self):
        # Create excel file with 2 sheets. Create a header in second sheet
        # where data will be inserted
        wb = openpyxl.Workbook()
        ws = wb.active
        ws.cell(row=1, column=1, value="My Test Value")
        ws2 = wb.create_sheet("data")
        ws2.cell(row=1, column=1, value='Partner Id')
        ws2.cell(row=1, column=2, value='Partner Name')
        output = BytesIO()
        wb.save(output)
        data = output.getvalue()

        # Create attachment with the created xlsx file which will be used as
        # template in the sql query
        attachmnent_vals = {
            'name': 'template xlsx sql export Res Partner',
            'datas': base64.b64encode(data),
        }
        attachment = self.env['ir.attachment'].create(attachmnent_vals)

        # Create the query and configure it to insert the data in the second
        # sheet of the xlsx template file and start inserting data at the
        # second row, ignoring header (because the template excel file
        # already contains a header)
        test_query = "SELECT id, name FROM res_partner"
        query_vals = {
            'name': 'Test Query Excel',
            'query': test_query,
            'file_format': 'excel',
            'attachment_id': attachment.id,
            'sheet_position': 2,
            'header': False,
            'row_position': 2,
        }
        query = self.env['sql.export'].create(query_vals)
        query.button_validate_sql_expression()
        wizard = self.wizard_obj.create({
            'sql_export_id': query.id,
        })

        # Check the generated excel file. The first sheet should still contain
        # the same data and the second sheet should have kept the header and
        # inserted data from the query
        wb2 = self.get_workbook_from_query(wizard)
        sheets = wb2.worksheets
        ws1 = sheets[0]
        # Check values, header should be here by default
        self.assertEqual(ws1.cell(row=1, column=1).value, 'My Test Value')
        ws2 = sheets[1]
        self.assertEqual(ws2.cell(row=1, column=1).value, 'Partner Id')
        self.assertTrue(ws2.cell(row=2, column=1).value)