#!/usr/bin/env python3
# group: rw quick
#
# Test for additional information emitted by qemu-img info on qcow2
# images
#
# Copyright (C) 2013 Red Hat, Inc.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

import os
import re
import json
import iotests
from iotests import qemu_img, qemu_img_info, supports_qcow2_zstd_compression
import unittest

test_img = os.path.join(iotests.test_dir, 'test.img')

class TestImageInfoSpecific(iotests.QMPTestCase):
    '''Abstract base class for ImageInfoSpecific tests'''

    def setUp(self):
        if self.img_options is None:
            self.skipTest('Skipping abstract test class')
        qemu_img('create', '-f', iotests.imgfmt, '-o', self.img_options,
                 test_img, '128K')

    def tearDown(self):
        os.remove(test_img)

class TestQemuImgInfo(TestImageInfoSpecific):
    '''Abstract base class for qemu-img info tests'''

    img_options = None
    json_compare = None
    human_compare = None

    def test_json(self):
        data = qemu_img_info(test_img)['format-specific']
        self.assertEqual(data['type'], iotests.imgfmt)
        self.assertEqual(data['data'], self.json_compare)

    def test_human(self):
        data = qemu_img('info', '--output=human', test_img).stdout.split('\n')
        data = data[(data.index('Format specific information:') + 1)
                    :data.index('')]
        for field in data:
            self.assertTrue(re.match('^ {4}[^ ]', field) is not None)
        data = [line.strip() for line in data]
        self.assertEqual(data, self.human_compare)

class TestQMP(TestImageInfoSpecific):
    '''Abstract base class for qemu QMP tests'''

    img_options = None
    qemu_options = ''
    TestImageInfoSpecific = TestImageInfoSpecific

    def setUp(self):
        self.TestImageInfoSpecific.setUp(self)
        self.vm = iotests.VM().add_drive(test_img, self.qemu_options)
        self.vm.launch()

    def tearDown(self):
        self.vm.shutdown()
        self.TestImageInfoSpecific.tearDown(self)

    def test_qmp(self):
        result = self.vm.qmp('query-block')['return']
        drive = next(drive for drive in result if drive['device'] == 'drive0')
        data = drive['inserted']['image']['format-specific']
        self.assertEqual(data['type'], iotests.imgfmt)
        self.assertEqual(data['data'], self.compare)

class TestQCow2(TestQemuImgInfo):
    '''Testing a qcow2 version 2 image'''
    img_options = 'compat=0.10,compression_type=zlib'
    json_compare = { 'compat': '0.10', 'refcount-bits': 16,
                     'compression-type': 'zlib' }
    human_compare = [ 'compat: 0.10', 'compression type: zlib',
                      'refcount bits: 16' ]

class TestQCow3NotLazy(TestQemuImgInfo):
    '''Testing a qcow2 version 3 image with lazy refcounts disabled'''
    if supports_qcow2_zstd_compression():
        compression_type = 'zstd'
    else:
        compression_type = 'zlib'

    img_options = 'compat=1.1,lazy_refcounts=off'
    img_options += f',compression_type={compression_type}'
    json_compare = { 'compat': '1.1', 'lazy-refcounts': False,
                     'refcount-bits': 16, 'corrupt': False,
                     'compression-type': compression_type, 'extended-l2': False }
    human_compare = [ 'compat: 1.1', f'compression type: {compression_type}',
                      'lazy refcounts: false', 'refcount bits: 16',
                      'corrupt: false', 'extended l2: false' ]

class TestQCow3Lazy(TestQemuImgInfo):
    '''Testing a qcow2 version 3 image with lazy refcounts enabled'''
    img_options = 'compat=1.1,lazy_refcounts=on,compression_type=zlib'
    json_compare = { 'compat': '1.1', 'lazy-refcounts': True,
                     'refcount-bits': 16, 'corrupt': False,
                     'compression-type': 'zlib', 'extended-l2': False }
    human_compare = [ 'compat: 1.1', 'compression type: zlib',
                      'lazy refcounts: true', 'refcount bits: 16',
                      'corrupt: false', 'extended l2: false' ]

class TestQCow3NotLazyQMP(TestQMP):
    '''Testing a qcow2 version 3 image with lazy refcounts disabled, opening
       with lazy refcounts enabled'''
    img_options = 'compat=1.1,lazy_refcounts=off,compression_type=zlib'
    qemu_options = 'lazy-refcounts=on'
    compare = { 'compat': '1.1', 'lazy-refcounts': False,
                'refcount-bits': 16, 'corrupt': False,
                'compression-type': 'zlib', 'extended-l2': False }


class TestQCow3LazyQMP(TestQMP):
    '''Testing a qcow2 version 3 image with lazy refcounts enabled, opening
       with lazy refcounts disabled'''
    if supports_qcow2_zstd_compression():
        compression_type = 'zstd'
    else:
        compression_type = 'zlib'

    img_options = 'compat=1.1,lazy_refcounts=on'
    img_options += f',compression_type={compression_type}'
    qemu_options = 'lazy-refcounts=off'
    compare = { 'compat': '1.1', 'lazy-refcounts': True,
                'refcount-bits': 16, 'corrupt': False,
                'compression-type': compression_type, 'extended-l2': False }

TestImageInfoSpecific = None
TestQemuImgInfo = None
TestQMP = None

if __name__ == '__main__':
    iotests.main(supported_fmts=['qcow2'],
                 supported_protocols=['file'],
                 unsupported_imgopts=['refcount_bits'])