#    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 mock

from nova import test
from nova.virt.xenapi import host


class VGPUTestCase(test.NoDBTestCase):
    """Unit tests for Driver operations."""
    @mock.patch.object(host.HostState, 'update_status',
                       return_value='fake_stats_1')
    @mock.patch.object(host.HostState, '_get_vgpu_stats_in_group')
    def test_get_vgpu_stats_empty_cfg(self, mock_get, mock_update):
        # no vGPU type configured.
        self.flags(enabled_vgpu_types=[], group='devices')
        session = mock.Mock()

        host_obj = host.HostState(session)
        stats = host_obj._get_vgpu_stats()

        session.call_xenapi.assert_not_called()
        self.assertEqual(stats, {})

    @mock.patch.object(host.HostState, 'update_status',
                       return_value='fake_stats_1')
    @mock.patch.object(host.HostState, '_get_vgpu_stats_in_group')
    def test_get_vgpu_stats_single_type(self, mock_get, mock_update):
        # configured single vGPU type
        self.flags(enabled_vgpu_types=['type_name_1'], group='devices')
        session = mock.Mock()
        # multiple GPU groups
        session.call_xenapi.side_effect = [
            ['grp_ref1', 'grp_ref2'],  # GPU_group.get_all
            'uuid_1',  # GPU_group.get_uuid
            'uuid_2',  # GPU_group.get_uuid
        ]
        # Let it return None for the 2nd GPU group for the case
        # that it doesn't have the specified vGPU type enabled.
        mock_get.side_effect = ['fake_stats_1', None]
        host_obj = host.HostState(session)
        stats = host_obj._get_vgpu_stats()

        self.assertEqual(session.call_xenapi.call_count, 3)
        self.assertEqual(mock_update.call_count, 1)
        self.assertEqual(mock_get.call_count, 2)
        self.assertEqual(stats, {'uuid_1': 'fake_stats_1'})

    @mock.patch.object(host.HostState, 'update_status',
                       return_value='fake_stats_1')
    @mock.patch.object(host.HostState, '_get_vgpu_stats_in_group')
    def test_get_vgpu_stats_multi_types(self, mock_get, mock_update):
        # when multiple vGPU types configured, it use the first one.
        self.flags(enabled_vgpu_types=['type_name_1', 'type_name_2'],
                   group='devices')
        session = mock.Mock()
        session.call_xenapi.side_effect = [
            ['grp_ref1'],  # GPU_group.get_all
            'uuid_1',  # GPU_group.get_uuid
        ]
        mock_get.side_effect = ['fake_stats_1']
        host_obj = host.HostState(session)
        stats = host_obj._get_vgpu_stats()

        self.assertEqual(session.call_xenapi.call_count, 2)
        self.assertEqual(mock_update.call_count, 1)
        self.assertEqual(stats, {'uuid_1': 'fake_stats_1'})
        # called with the first vGPU type: 'type_name_1'
        mock_get.assert_called_with('grp_ref1', ['type_name_1'])

    @mock.patch.object(host.HostState, 'update_status',
                       return_value='fake_stats_1')
    @mock.patch.object(host.HostState, '_get_total_vgpu_in_grp',
                       return_value=7)
    def test_get_vgpu_stats_in_group(self, mock_get, mock_update):
        # Test it will return vGPU stat for the enabled vGPU type.
        enabled_vgpu_types = ['type_name_2']
        session = mock.Mock()
        session.call_xenapi.side_effect = [
            ['type_ref_1', 'type_ref_2'],  # GPU_group.get_enabled_VGPU_types
            'type_name_1',  # VGPU_type.get_model_name
            'type_name_2',  # VGPU_type.get_model_name