# Copyright 2012 Michael Still and Canonical Inc
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.


import os
import tempfile
import time
from unittest import mock

import eventlet
import fixtures

from nova import test
from nova.virt.disk.mount import nbd
from nova.virt.image import model as imgmodel

ORIG_EXISTS = os.path.exists
ORIG_LISTDIR = os.listdir


def _fake_exists_no_users(path):
    if path.startswith('/sys/block/nbd'):
        if path.endswith('pid'):
            return False
        return True
    return ORIG_EXISTS(path)


def _fake_listdir_nbd_devices(path):
    if isinstance(path, str) and path.startswith('/sys/block'):
        return ['nbd0', 'nbd1']
    return ORIG_LISTDIR(path)


def _fake_exists_all_used(path):
    if path.startswith('/sys/block/nbd'):
        return True
    return ORIG_EXISTS(path)


def _fake_detect_nbd_devices_none():
    return []


def _fake_detect_nbd_devices():
    return ['nbd0', 'nbd1']


def _fake_noop(*args, **kwargs):
    return


class NbdTestCaseNoStub(test.NoDBTestCase):

    @mock.patch('os.listdir')
    def test_detect_nbd_devices(self, list_dir_mock):
        list_dir_mock.return_value = _fake_detect_nbd_devices()
        result = nbd.NbdMount._detect_nbd_devices()
        self.assertIsNotNone(result)
        self.assertIsInstance(result, list)
        self.assertEqual(len(list_dir_mock.return_value), len(result))
        for path in list_dir_mock.return_value:
            self.assertIn(path, result)

    @mock.patch('os.listdir')
    def test_detect_nbd_devices_empty(self, list_dir_mock):
        list_dir_mock.return_value = [
            "nbdz", "fake0", "not-nbd1"]
        result = nbd.NbdMount._detect_nbd_devices()
        self.assertIsNotNone(result)
        self.assertIsInstance(result, list)
        self.assertEqual(0, len(result))


class NbdTestCase(test.NoDBTestCase):
    def setUp(self):
        super(NbdTestCase, self).setUp()
        self.stub_out('nova.virt.disk.mount.nbd.NbdMount._detect_nbd_devices',
                      _fake_detect_nbd_devices)
        self.useFixture(fixtures.MonkeyPatch('os.listdir',
                                             _fake_listdir_nbd_devices))
        self.file = imgmodel.LocalFileImage("/some/file.qcow2",
                                            imgmodel.FORMAT_QCOW2)

    def test_nbd_no_devices(self):
        tempdir = self.useFixture(fixtures.TempDir()).path
        self.stub_out('nova.virt.disk.mount.nbd.NbdMount._detect_nbd_devices',
                      _fake_detect_nbd_devices_none)
        n = nbd.NbdMount(self.file, tempdir)
        self.assertIsNone(n._allocate_nbd())

    def test_nbd_no_free_devices(self):
        tempdir = self.useFixture(fixtures.TempDir()).path
        n = nbd.NbdMount(self.file, tempdir)
        self.useFixture(fixtures.MonkeyPatch('os.path.exists',
                                             _fake_exists_all_used))
        self.assertIsNone(n._allocate_nbd())

    def test_nbd_not_loaded(self):
        tempdir = self.useFixture(fixtures.TempDir()).path
        n = nbd.NbdMount(self.file, tempdir)

        # Fake out os.path.exists
        def fake_exists(path):
            if path.startswith('/sys/block/nbd'):
                return False
            return ORIG_EXISTS(path)
        self.useFixture(fixtures.MonkeyPatch('os.path.exists', fake_exists))

        # This should fail, as we don't have the module "loaded"
        # TODO(mikal): work out how to force english as the gettext language
        # so that the error check always passes
        self.assertIsNone(n._allocate_nbd())
        self.assertEqual('nbd unavailable: module not loaded', n.error)

    def test_nbd_allocation(self):
        tempdir = self.useFixture(fixtures.TempDir()).path
        n = nbd.NbdMount(self.file, tempdir)
        self.useFixture(fixtures.MonkeyPatch('os.path.exists',
                                             _fake_exists_no_users))
        self.useFixture(fixtures.MonkeyPatch('random.shuffle', _fake_noop))

        # Allocate a nbd device
        self.assertEqual('/dev/nbd0', n._allocate_nbd())

    def test_nbd_allocation_one_in_use(self):
        tempdir = self.useFixture(fixtures.TempDir()).path
        n = nbd.NbdMount(self.file, tempdir)
        self.useFixture(fixtures.MonkeyPatch('random.shuffle', _fake_noop))

        # Fake out os.path.exists
        def fake_exists(path):
            if path.startswith('/sys/block/nbd'):
                if path == '/sys/block/nbd0/pid':
                    return True
                if path.endswith('pid'):
                    return False
                return True
            return ORIG_EXISTS(path)
        self.useFixture(fixtures.MonkeyPatch('os.path.exists', fake_exists))

        # Allocate a nbd device, should not be the in use one
        # TODO(mikal): Note that there is a leak here, as the in use nbd device
        # is removed from the list, but not returned so it will never be
        # re-added. I will fix this in a later patch.
        self.assertEqual('/dev/nbd1', n._allocate_nbd())

    def test_inner_get_dev_no_devices(self):
        tempdir = self.useFixture(fixtures.TempDir()).path
        self.stub_out('nova.virt.disk.mount.nbd.NbdMount._detect_nbd_devices',
                      _fake_detect_nbd_devices_none)
        n = nbd.NbdMount(self.file, tempdir)
        self.assertFalse(n._inner_get_dev())

    @mock.patch('nova.privsep.fs.nbd_connect', return_value=('', 'broken'))
    def test_inner_get_dev_qemu_fails(self, mock_nbd_connect):
        tempdir = self.useFixture(fixtures.TempDir()).path
        n = nbd.NbdMount(self.file, tempdir)
        self.useFixture(fixtures.MonkeyPatch('os.path.exists',
                                             _fake_exists_no_users))

        # Error logged, no device consumed
        self.assertFalse(n._inner_get_dev())
        self.assertTrue(n.error.startswith('qemu-nbd error'))

    @mock.patch('random.shuffle')
    @mock.patch('os.listdir', return_value=['nbd0', 'nbd1', 'loop0'])
    @mock.patch('nova.privsep.fs.nbd_connect', return_value=('', ''))
    @mock.patch('nova.privsep.fs.nbd_disconnect', return_value=('', ''))
    @mock.patch('time.sleep')
    def test_inner_get_dev_qemu_timeout(self, mock_sleep, mock_nbd_disconnct,
                                        mock_nbd_connect, mock_listdir,
                                        mock_shuffle):
        self.flags(timeout_nbd=3)
        tempdir = self.useFixture(fixtures.TempDir()).path

        with mock.patch('os.path.exists', side_effect=[
                    True, False, False, False, False, False, False, False]):
            n = nbd.NbdMount(self.file, tempdir)

            # Error logged, no device consumed
            self.assertFalse(n._inner_get_dev())
            self.assertTrue(n.error.endswith('did not show up'))

    @mock.patch('time.sleep', new=mock.Mock())
    @mock.patch('random.shuffle')
    @mock.patch('os.path.exists', side_effect=[True, False, False, False,
                                               False, True])
    @mock.patch('os.listdir', return_value=['nbd0', 'nbd1', 'loop0'])
    @mock.patch('nova.privsep.fs.nbd_connect', return_value=('', ''))
    @mock.patch('nova.privsep.fs.nbd_disconnect')
    def test_inner_get_dev_works(self, mock_nbd_disconnect, mock_nbd_connect,
                                 mock_exists, mock_listdir, mock_shuffle):
        tempdir = self.useFixture(fixtures.TempDir()).path
        n = nbd.NbdMount(self.file, tempdir)

        # No error logged, device consumed
        self.assertTrue(n._inner_get_dev())
        self.assertTrue(n.linked)
        self.assertEqual('', n.error)
        self.assertEqual('/dev/nbd0', n.device)

        # Free
        n.unget_dev()
        self.assertFalse(n.linked)
        self.assertEqual('', n.error)
        self.assertIsNone(n.device)

    def test_unget_dev_simple(self):
        # This test is just checking we don't get an exception when we unget
        # something we don't have
        tempdir = self.useFixture(fixtures.TempDir()).path
        n = nbd.NbdMount(self.file, tempdir)
        n.unget_dev()

    @mock.patch('time.sleep', new=mock.Mock())
    @mock.patch('random.shuffle')
    @mock.patch('os.path.exists', side_effect=[True, False, False, False,
                                               False, True])
    @mock.patch('os.listdir', return_value=['nbd0', 'nbd1', 'loop0'])
    @mock.patch('nova.privsep.fs.nbd_connect', return_value=('', ''))
    @mock.patch('nova.privsep.fs.nbd_disconnect')
    def test_get_dev(self, mock_nbd_disconnect, mock_nbd_connect,
                     mock_exists, mock_listdir, mock_shuffle):
        tempdir = self.useFixture(fixtures.TempDir()).path
        n = nbd.NbdMount(self.file, tempdir)

        # No error logged, device consumed
        self.assertTrue(n.get_dev())
        self.assertTrue(n.linked)
        self.assertEqual('', n.error)
        self.assertEqual('/dev/nbd0', n.device)

        # Free
        n.unget_dev()
        self.assertFalse(n.linked)
        self.assertEqual('', n.error)
        self.assertIsNone(n.device)

    @mock.patch('random.shuffle')
    @mock.patch('time.sleep')
    @mock.patch('nova.privsep.fs.nbd_connect')
    @mock.patch('nova.privsep.fs.nbd_disconnect')
    @mock.patch('os.path.exists', return_value=True)
    @mock.patch('nova.virt.disk.mount.nbd.NbdMount._inner_get_dev',
                return_value=False)
    def test_get_dev_timeout(self, mock_get_dev, mock_exists,
                             mock_nbd_disconnect, mock_nbd_connect,
                             mock_sleep, mock_shuffle):
        tempdir = self.useFixture(fixtures.TempDir()).path
        n = nbd.NbdMount(self.file, tempdir)
        self.useFixture(fixtures.MonkeyPatch(('nova.virt.disk.mount.api.'
                                              'MAX_DEVICE_WAIT'), -10))

        # No error logged, device consumed
        self.assertFalse(n.get_dev())

    @mock.patch('nova.privsep.fs.mount', return_value=('', 'broken'))
    def test_do_mount_need_to_specify_fs_type(self, mock_mount):
        # NOTE(mikal): Bug 1094373 saw a regression where we failed to
        # communicate a failed mount properly.
        imgfile = tempfile.NamedTemporaryFile()
        self.addCleanup(imgfile.close)
        tempdir = self.useFixture(fixtures.TempDir()).path
        mount = nbd.NbdMount(imgfile.name, tempdir)

        def fake_returns_true(*args, **kwargs):
            return True
        mount.get_dev = fake_returns_true
        mount.map_dev = fake_returns_true

        self.assertFalse(mount.do_mount())
        mock_mount.assert_called_once_with(
                None, None, tempdir, None)

    @mock.patch('nova.privsep.fs.nbd_connect')
    @mock.patch('nova.privsep.fs.nbd_disconnect')
    @mock.patch('os.path.exists')
    def test_device_creation_race(self, mock_exists, mock_nbd_disconnect,
                                  mock_nbd_connect):
        # Make sure that even if two threads create instances at the same time
        # they cannot choose the same nbd number (see bug 1207422)

        tempdir = self.useFixture(fixtures.TempDir()).path
        free_devices = _fake_detect_nbd_devices()[:]
        chosen_devices = []

        def fake_find_unused(self):
            return os.path.join('/dev', free_devices[-1])

        def delay_and_remove_device(*args, **kwargs):
            # Ensure that context switch happens before the device is marked
            # as used. This will cause a failure without nbd-allocation-lock
            # in place.
            time.sleep(0.1)

            # We always choose the top device in find_unused - remove it now.
            free_devices.pop()

            return '', ''

        def pid_exists(pidfile):
            return pidfile not in [os.path.join('/sys/block', dev, 'pid')
                                   for dev in free_devices]

        self.stub_out('nova.virt.disk.mount.nbd.NbdMount._allocate_nbd',
                      fake_find_unused)
        mock_nbd_connect.side_effect = delay_and_remove_device
        mock_exists.side_effect = pid_exists

        def get_a_device():
            n = nbd.NbdMount(self.file, tempdir)
            n.get_dev()
            chosen_devices.append(n.device)

        thread1 = eventlet.spawn(get_a_device)
        thread2 = eventlet.spawn(get_a_device)
        thread1.wait()
        thread2.wait()

        self.assertEqual(2, len(chosen_devices))
        self.assertNotEqual(chosen_devices[0], chosen_devices[1])
