#!/usr/bin/python3
#
# Tests for inserting and removing filters in a block graph.
#
# Copyright (c) 2022 Virtuozzo International GmbH.
#
# 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 iotests
from iotests import qemu_img_create, try_remove


disk = os.path.join(iotests.test_dir, 'disk')
sock = os.path.join(iotests.sock_dir, 'sock')
size = 1024 * 1024


class TestFilterInsertion(iotests.QMPTestCase):
    def setUp(self):
        qemu_img_create(disk, str(size))

        self.vm = iotests.VM()
        self.vm.set_machine('pc')  # to make device_add work
        self.vm.launch()

        self.vm.qmp_check('blockdev-add', {
            'node-name': 'disk0',
            'driver': 'qcow2',
            'file': {
                'node-name': 'file0',
                'driver': 'file',
                'filename': disk
            }
        })

    def tearDown(self):
        self.vm.shutdown()
        os.remove(disk)
        try_remove(sock)

    def test_simple_insertion(self):
        vm = self.vm

        vm.qmp_check('blockdev-add', {
            'node-name': 'filter',
            'driver': 'preallocate',
            'file': 'file0'
        })

        vm.qmp_check('vz-x-blockdev-replace', {
            'parent-type': 'driver',
            'node-name': 'disk0',
            'child': 'file',
            'new-child': 'filter'
        })

        # Filter inserted:
        # disk0 -file-> filter -file-> file0
        vm.assert_edges_list([
            ('disk0', 'file', 'filter'),
            ('filter', 'file', 'file0')
        ])

        vm.qmp_check('vz-x-blockdev-replace', {
            'parent-type': 'driver',
            'node-name': 'disk0',
            'child': 'file',
            'new-child': 'file0'
        })

        # Filter replaced, but still exists:
        # dik0 -file-> file0 <-file- filter
        vm.assert_edges_list([
            ('disk0', 'file', 'file0'),
            ('filter', 'file', 'file0')
        ])

        vm.qmp_check('blockdev-del', node_name='filter')

        # Filter removed
        # dik0 -file-> file0
        vm.assert_edges_list([
            ('disk0', 'file', 'file0')
        ])

    def test_tran_insert_under_qdev(self):
        vm = self.vm

        vm.qmp_check('device_add', driver='virtio-scsi')
        vm.qmp_check('device_add', id='sda', driver='scsi-hd', drive='disk0')

        vm.qmp_check('transaction', actions=[
            {
                'type': 'vz-x-blockdev-add',
                'data': {
                    'node-name': 'filter',
                    'driver': 'compress',
                    'file': 'disk0'
                }
            }, {
                'type': 'vz-x-blockdev-replace',
                'data': {
                    'parent-type': 'qdev',
                    'qdev-id': 'sda',
                    'new-child': 'filter'
                }
            }
        ])

        # Filter inserted:
        # sda -root-> filter -file-> disk0 -file-> file0
        vm.assert_edges_list([
            # parent_node_name, child_name, child_node_name
            ('sda', 'root', 'filter'),
            ('filter', 'file', 'disk0'),
            ('disk0', 'file', 'file0'),
        ])

        vm.qmp_check('vz-x-blockdev-replace', {
            'parent-type': 'qdev',
            'qdev-id': 'sda',
            'new-child': 'disk0'
        })
        vm.qmp_check('blockdev-del', node_name='filter')

        # Filter removed:
        # sda -root-> disk0 -file-> file0
        vm.assert_edges_list([
            # parent_node_name, child_name, child_node_name
            ('sda', 'root', 'disk0'),
            ('disk0', 'file', 'file0'),
        ])

    def test_tran_insert_under_nbd_export(self):
        vm = self.vm

        vm.qmp_check('nbd-server-start',
                     addr={'type': 'unix', 'data': {'path': sock}})
        vm.qmp_check('block-export-add', id='exp1', type='nbd',
                     node_name='disk0', name='exp1')
        vm.qmp_check('block-export-add', id='exp2', type='nbd',
                     node_name='disk0', name='exp2')
        vm.qmp_check('object-add', qom_type='throttle-group',
                     id='tg', limits={'iops-read': 1})

        vm.qmp_check('transaction', actions=[
            {
                'type': 'vz-x-blockdev-add',
                'data': {
                    'node-name': 'filter',
                    'driver': 'throttle',
                    'throttle-group': 'tg',
                    'file': 'disk0'
                }
            }, {
                'type': 'vz-x-blockdev-replace',
                'data': {
                    'parent-type': 'export',
                    'export-id': 'exp1',
                    'new-child': 'filter'
                }
            }
        ])

        # Only exp1 is throttled, exp2 is not:
        # exp1 -root-> filter
        #                |
        #                |file
        #                v
        # exp2 -file-> disk0 -file> file0
        vm.assert_edges_list([
            # parent_node_name, child_name, child_node_name
            ('exp1', 'root', 'filter'),
            ('filter', 'file', 'disk0'),
            ('disk0', 'file', 'file0'),
            ('exp2', 'root', 'disk0')
        ])

        vm.qmp_check('vz-x-blockdev-replace', {
            'parent-type': 'export',
            'export-id': 'exp2',
            'new-child': 'filter'
        })

        # Both throttled:
        # exp1 -root-> filter <-file- exp2
        #                |
        #                |file
        #                v
        #              disk0 -file> file0
        vm.assert_edges_list([
            # parent_node_name, child_name, child_node_name
            ('exp1', 'root', 'filter'),
            ('filter', 'file', 'disk0'),
            ('disk0', 'file', 'file0'),
            ('exp2', 'root', 'filter')
        ])

        # Check, that filter is in use and can't be removed
        result = vm.qmp('blockdev-del', node_name='filter')
        self.assert_qmp(result, 'error/desc', 'Node filter is in use')

        vm.qmp_check('transaction', actions=[
            {
                'type': 'vz-x-blockdev-replace',
                'data': {
                    'parent-type': 'export',
                    'export-id': 'exp1',
                    'new-child': 'disk0'
                }
            }, {
                'type': 'vz-x-blockdev-replace',
                'data': {
                    'parent-type': 'export',
                    'export-id': 'exp2',
                    'new-child': 'disk0'
                }
            }
        ])
        vm.qmp_check('blockdev-del', node_name='filter')

        # Filter removed:
        # exp1 -root-> disk0 <-file- exp2
        #                |
        #                |file
        #                v
        #              file0
        vm.assert_edges_list([
            # parent_node_name, child_name, child_node_name
            ('exp1', 'root', 'disk0'),
            ('disk0', 'file', 'file0'),
            ('exp2', 'root', 'disk0')
        ])


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