// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.
//

package com.cloud.template;


import com.cloud.exception.InvalidParameterValueException;
import com.cloud.exception.ResourceAllocationException;
import com.cloud.host.Status;
import com.cloud.hypervisor.Hypervisor;
import com.cloud.hypervisor.HypervisorGuruManager;
import com.cloud.resourcelimit.CheckedReservation;
import com.cloud.storage.DataStoreRole;
import com.cloud.storage.GuestOSVO;
import com.cloud.storage.Snapshot;
import com.cloud.storage.SnapshotVO;
import com.cloud.storage.Storage;
import com.cloud.storage.StoragePool;
import com.cloud.storage.StoragePoolStatus;
import com.cloud.storage.TemplateProfile;
import com.cloud.storage.VMTemplateStoragePoolVO;
import com.cloud.storage.VMTemplateStorageResourceAssoc;
import com.cloud.storage.VMTemplateVO;
import com.cloud.storage.VolumeVO;
import com.cloud.storage.dao.GuestOSDao;
import com.cloud.storage.dao.SnapshotDao;
import com.cloud.storage.dao.StoragePoolHostDao;
import com.cloud.storage.dao.VMTemplateDao;
import com.cloud.storage.dao.VMTemplateDetailsDao;
import com.cloud.storage.dao.VMTemplatePoolDao;
import com.cloud.storage.dao.VolumeDao;
import com.cloud.user.Account;
import com.cloud.user.AccountManager;
import com.cloud.user.AccountVO;
import com.cloud.user.ResourceLimitService;
import com.cloud.user.User;
import com.cloud.user.UserData;
import com.cloud.user.UserDataVO;
import com.cloud.user.UserVO;
import com.cloud.user.dao.UserDataDao;
import com.cloud.utils.UriUtils;
import com.cloud.utils.concurrency.NamedThreadFactory;
import com.cloud.utils.exception.CloudRuntimeException;
import com.cloud.vm.VMInstanceVO;
import com.cloud.vm.dao.VMInstanceDao;
import junit.framework.TestCase;

import org.apache.cloudstack.api.command.user.iso.RegisterIsoCmd;
import org.apache.cloudstack.api.command.user.template.CreateTemplateCmd;
import org.apache.cloudstack.api.command.user.template.DeleteTemplateCmd;
import org.apache.cloudstack.api.command.user.template.RegisterTemplateCmd;
import org.apache.cloudstack.api.command.user.template.RegisterVnfTemplateCmd;
import org.apache.cloudstack.api.command.user.template.UpdateTemplateCmd;
import org.apache.cloudstack.api.command.user.template.UpdateVnfTemplateCmd;
import org.apache.cloudstack.api.command.user.userdata.LinkUserDataToTemplateCmd;
import org.apache.cloudstack.context.CallContext;
import org.apache.cloudstack.engine.subsystem.api.storage.DataStore;
import org.apache.cloudstack.engine.subsystem.api.storage.DataStoreManager;
import org.apache.cloudstack.engine.subsystem.api.storage.PrimaryDataStore;
import org.apache.cloudstack.engine.subsystem.api.storage.StorageStrategyFactory;
import org.apache.cloudstack.reservation.dao.ReservationDao;
import org.apache.cloudstack.secstorage.heuristics.HeuristicType;
import org.apache.cloudstack.storage.datastore.db.ImageStoreDao;
import org.apache.cloudstack.storage.datastore.db.ImageStoreVO;
import org.apache.cloudstack.storage.datastore.db.PrimaryDataStoreDao;
import org.apache.cloudstack.storage.datastore.db.StoragePoolVO;
import org.apache.cloudstack.storage.datastore.db.TemplateDataStoreDao;
import org.apache.cloudstack.storage.datastore.db.TemplateDataStoreVO;
import org.apache.cloudstack.storage.heuristics.HeuristicRuleHelper;
import org.apache.cloudstack.storage.template.VnfTemplateManager;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.MockedConstruction;
import org.mockito.MockedStatic;
import org.mockito.Mockito;
import org.mockito.Spy;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.junit.MockitoJUnitRunner;
import org.mockito.stubbing.Answer;

import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyBoolean;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.nullable;
import static org.mockito.Mockito.doNothing;
import static org.mockito.Mockito.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

@RunWith(MockitoJUnitRunner.class)
public class TemplateManagerImplTest extends TestCase {

    @Spy
    @InjectMocks
    TemplateManagerImpl templateManager;

    @Mock
    DataStoreManager dataStoreManager;

    @Mock
    VMTemplateDao vmTemplateDao;

    @Mock
    VMTemplatePoolDao vmTemplatePoolDao;

    @Mock
    TemplateDataStoreDao templateDataStoreDao;

    @Mock
    StoragePoolHostDao storagePoolHostDao;

    @Mock
    PrimaryDataStoreDao primaryDataStoreDao;

    @Mock
    ResourceLimitService resourceLimitMgr;

    @Mock
    ImageStoreDao imgStoreDao;

    @Mock
    GuestOSDao guestOSDao;

    @Mock
    SnapshotDao snapshotDao;

    @Mock
    VolumeDao volumeDao;

    @Mock
    VMTemplateDetailsDao tmpltDetailsDao;

    @Mock
    StorageStrategyFactory storageStrategyFactory;

    @Mock
    VMInstanceDao _vmInstanceDao;

    @Mock
    ReservationDao reservationDao;

    @Mock
    HypervisorGuruManager _hvGuruMgr;

    @Mock
    AccountManager _accountMgr;

    @Mock
    VnfTemplateManager vnfTemplateManager;

    @Mock
    HeuristicRuleHelper heuristicRuleHelperMock;

    @Mock
    private UserDataDao userDataDaoMock;

    @Mock
    private UserDataVO userDataMock;

    public class CustomThreadPoolExecutor extends ThreadPoolExecutor {
        AtomicInteger ai = new AtomicInteger(0);
        public CustomThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit,
                                        BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
            super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
        }

        @Override
        protected void beforeExecute(Thread t, Runnable r) {
            ai.addAndGet(1);
        }

        public int getCount() {
            try {
                // Wait for some time to give before execute to run. Otherwise the tests that
                // assert and check that template seeding has been scheduled may fail. If tests
                // are seen to fail, consider increasing the sleep time.
                Thread.sleep(1000);
                return ai.get();
            } catch (Exception e) {
                return -1;
            }
        }
    }

    @Before
    public void setUp() {
        AccountVO account = new AccountVO("admin", 1L, "networkDomain", Account.Type.NORMAL, "uuid");
        UserVO user = new UserVO(1, "testuser", "password", "firstname", "lastName", "email", "timezone", UUID.randomUUID().toString(), User.Source.UNKNOWN);
        CallContext.register(user, account);
    }

    @After
    public void tearDown() {
        CallContext.unregister();
    }

    @Test(expected = InvalidParameterValueException.class)
    public void testVerifyTemplateIdOfSystemTemplate() {
        templateManager.verifyTemplateId(1L);
    }

    public void testVerifyTemplateIdOfNonSystemTemplate() {
        templateManager.verifyTemplateId(1L);
    }

    @Test
    public void testForceDeleteTemplate() {
        //In this Unit test all the conditions related to "force template delete flag" are tested.

        DeleteTemplateCmd cmd = mock(DeleteTemplateCmd.class);
        VMTemplateVO template = mock(VMTemplateVO.class);
        TemplateAdapter templateAdapter = mock(TemplateAdapter.class);
        TemplateProfile templateProfile = mock(TemplateProfile.class);


        List<VMInstanceVO> vmInstanceVOList  = new ArrayList<VMInstanceVO>();
        List<TemplateAdapter> adapters  = new ArrayList<TemplateAdapter>();
        adapters.add(templateAdapter);
        when(cmd.getId()).thenReturn(0L);
        when(vmTemplateDao.findById(cmd.getId())).thenReturn(template);
        when(cmd.getZoneId()).thenReturn(null);

        when(template.getHypervisorType()).thenReturn(Hypervisor.HypervisorType.None);
        when(template.getFormat()).thenReturn(Storage.ImageFormat.VMDK);
        templateManager.setTemplateAdapters(adapters);
        when(templateAdapter.getName()).thenReturn(TemplateAdapter.TemplateAdapterType.Hypervisor.getName().toString());
        when(templateAdapter.prepareDelete(any(DeleteTemplateCmd.class))).thenReturn(templateProfile);
        when(templateAdapter.delete(templateProfile)).thenReturn(true);

        //case 1: When Force delete flag is 'true' but VM instance VO list is empty.
        when(cmd.isForced()).thenReturn(true);
        templateManager.deleteTemplate(cmd);

        //case 2.1: When Force delete flag is 'false' and VM instance VO list is empty.
        when(cmd.isForced()).thenReturn(false);
        templateManager.deleteTemplate(cmd);

        //case 2.2: When Force delete flag is 'false' and VM instance VO list is non empty.
        when(cmd.isForced()).thenReturn(false);
        VMInstanceVO vmInstanceVO = mock(VMInstanceVO.class);
        vmInstanceVOList.add(vmInstanceVO);
        when(_vmInstanceDao.listNonExpungedByTemplate(anyLong())).thenReturn(vmInstanceVOList);
        try {
            templateManager.deleteTemplate(cmd);
        } catch(Exception e) {
            assertTrue("Invalid Parameter Value Exception is expected", (e instanceof InvalidParameterValueException));
        }
    }
    @Test
    public void testPrepareTemplateIsSeeded() {
        VMTemplateVO mockTemplate = mock(VMTemplateVO.class);
        when(mockTemplate.getId()).thenReturn(202l);

        StoragePoolVO mockPool = mock(StoragePoolVO.class);

        PrimaryDataStore mockPrimaryDataStore = mock(PrimaryDataStore.class);
        when(mockPrimaryDataStore.getId()).thenReturn(2l);

        VMTemplateStoragePoolVO mockTemplateStore = mock(VMTemplateStoragePoolVO.class);
        when(mockTemplateStore.getDownloadState()).thenReturn(VMTemplateStorageResourceAssoc.Status.DOWNLOADED);

        when(vmTemplateDao.findById(anyLong(), anyBoolean())).thenReturn(mockTemplate);
        when(vmTemplatePoolDao.findByPoolTemplate(anyLong(), anyLong(), nullable(String.class))).thenReturn(mockTemplateStore);

        doNothing().when(mockTemplateStore).setMarkedForGC(anyBoolean());

        VMTemplateStoragePoolVO returnObject = templateManager.prepareTemplateForCreate(mockTemplate, (StoragePool) mockPrimaryDataStore);
        assertTrue("Test template is already seeded", returnObject == mockTemplateStore);
    }

    @Test
    public void testPrepareTemplateNotDownloaded() {
        VMTemplateVO mockTemplate = mock(VMTemplateVO.class);
        when(mockTemplate.getId()).thenReturn(202l);

        StoragePoolVO mockPool = mock(StoragePoolVO.class);

        PrimaryDataStore mockPrimaryDataStore = mock(PrimaryDataStore.class);
        when(mockPrimaryDataStore.getId()).thenReturn(2l);
        when(mockPrimaryDataStore.getDataCenterId()).thenReturn(1l);

        when(vmTemplateDao.findById(anyLong(), anyBoolean())).thenReturn(mockTemplate);
        when(vmTemplatePoolDao.findByPoolTemplate(anyLong(), anyLong(), nullable(String.class))).thenReturn(null);
        when(templateDataStoreDao.findByTemplateZoneDownloadStatus(202l, 1l, VMTemplateStorageResourceAssoc.Status.DOWNLOADED)).thenReturn(null);

        VMTemplateStoragePoolVO returnObject = templateManager.prepareTemplateForCreate(mockTemplate, (StoragePool) mockPrimaryDataStore);
        assertTrue("Test template is not ready", returnObject == null);
    }

    @Test(expected = CloudRuntimeException.class)
    public void testPrepareTemplateNoHostConnectedToPool() {
        VMTemplateVO mockTemplate = mock(VMTemplateVO.class);
        when(mockTemplate.getId()).thenReturn(202l);

        StoragePoolVO mockPool = mock(StoragePoolVO.class);

        PrimaryDataStore mockPrimaryDataStore = mock(PrimaryDataStore.class);
        when(mockPrimaryDataStore.getId()).thenReturn(2l);
        when(mockPrimaryDataStore.getDataCenterId()).thenReturn(1l);

        TemplateDataStoreVO mockTemplateDataStore = mock(TemplateDataStoreVO.class);

        when(vmTemplateDao.findById(anyLong(), anyBoolean())).thenReturn(mockTemplate);
        when(vmTemplatePoolDao.findByPoolTemplate(anyLong(), anyLong(), nullable(String.class))).thenReturn(null);
        when(templateDataStoreDao.findByTemplateZoneDownloadStatus(202l, 1l, VMTemplateStorageResourceAssoc.Status.DOWNLOADED)).thenReturn(mockTemplateDataStore);
        when(storagePoolHostDao.listByHostStatus(2l, Status.Up)).thenReturn(null);

        templateManager.prepareTemplateForCreate(mockTemplate, (StoragePool) mockPrimaryDataStore);
    }

    @Test(expected = InvalidParameterValueException.class)
    public void testPrepareTemplateInvalidTemplate() {
        when(vmTemplateDao.findById(anyLong())).thenReturn(null);
        templateManager.prepareTemplate(202, 1, null);
    }

    @Test
    public void testTemplateScheduledForDownloadInOnePool() {
        VMTemplateVO mockTemplate = mock(VMTemplateVO.class);
        StoragePoolVO mockPool = mock(StoragePoolVO.class);
        PrimaryDataStore mockPrimaryDataStore = mock(PrimaryDataStore.class);
        VMTemplateStoragePoolVO mockTemplateStore = mock(VMTemplateStoragePoolVO.class);

        when(mockPrimaryDataStore.getId()).thenReturn(2l);
        when(mockPool.getId()).thenReturn(2l);
        when(mockPool.getStatus()).thenReturn(StoragePoolStatus.Up);
        when(mockPool.getDataCenterId()).thenReturn(1l);
        when(mockTemplate.getId()).thenReturn(202l);
        when(mockTemplateStore.getDownloadState()).thenReturn(VMTemplateStorageResourceAssoc.Status.DOWNLOADED);
        when(vmTemplateDao.findById(anyLong())).thenReturn(mockTemplate);
        when(dataStoreManager.getPrimaryDataStore(anyLong())).thenReturn(mockPrimaryDataStore);
        when(vmTemplateDao.findById(anyLong(), anyBoolean())).thenReturn(mockTemplate);
        when(vmTemplatePoolDao.findByPoolTemplate(anyLong(), anyLong(), nullable(String.class))).thenReturn(mockTemplateStore);
        when(primaryDataStoreDao.findById(anyLong())).thenReturn(mockPool);

        doNothing().when(mockTemplateStore).setMarkedForGC(anyBoolean());

        ExecutorService preloadExecutor = new CustomThreadPoolExecutor(8, 8, 0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue(),
                new NamedThreadFactory("Template-Preloader"));
        templateManager._preloadExecutor = preloadExecutor;

        templateManager.prepareTemplate(202, 1, 2l);
        assertTrue("Test template is scheduled for seeding to on pool", ((CustomThreadPoolExecutor)preloadExecutor).getCount() == 1);
    }

    @Test
    public void testTemplateScheduledForDownloadInDisabledPool() {
        VMTemplateVO mockTemplate = mock(VMTemplateVO.class);
        StoragePoolVO mockPool = mock(StoragePoolVO.class);
        PrimaryDataStore mockPrimaryDataStore = mock(PrimaryDataStore.class);
        VMTemplateStoragePoolVO mockTemplateStore = mock(VMTemplateStoragePoolVO.class);

        when(mockPool.getStatus()).thenReturn(StoragePoolStatus.Disabled);
        when(vmTemplateDao.findById(anyLong())).thenReturn(mockTemplate);
        when(primaryDataStoreDao.findById(anyLong())).thenReturn(mockPool);

        ExecutorService preloadExecutor = new CustomThreadPoolExecutor(8, 8, 0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue(),
                new NamedThreadFactory("Template-Preloader"));
        templateManager._preloadExecutor = preloadExecutor;

        templateManager.prepareTemplate(202, 1, 2l);
        assertTrue("Test template is not scheduled for seeding on disabled pool", ((CustomThreadPoolExecutor)preloadExecutor).getCount() == 0);
    }

    @Test
    public void testTemplateScheduledForDownloadInMultiplePool() {
        VMTemplateVO mockTemplate = mock(VMTemplateVO.class);
        PrimaryDataStore mockPrimaryDataStore = mock(PrimaryDataStore.class);
        VMTemplateStoragePoolVO mockTemplateStore = mock(VMTemplateStoragePoolVO.class);
        List<StoragePoolVO> pools = new ArrayList<StoragePoolVO>();

        StoragePoolVO mockPool1 = mock(StoragePoolVO.class);
        when(mockPool1.getId()).thenReturn(2l);
        when(mockPool1.getDataCenterId()).thenReturn(1l);
        StoragePoolVO mockPool2 = mock(StoragePoolVO.class);
        when(mockPool2.getDataCenterId()).thenReturn(1l);
        StoragePoolVO mockPool3 = mock(StoragePoolVO.class);
        when(mockPool3.getDataCenterId()).thenReturn(2l);
        pools.add(mockPool1);
        pools.add(mockPool2);
        pools.add(mockPool3);

        when(mockPrimaryDataStore.getId()).thenReturn(2l);
        when(mockTemplate.getId()).thenReturn(202l);
        when(mockTemplateStore.getDownloadState()).thenReturn(VMTemplateStorageResourceAssoc.Status.DOWNLOADED);
        when(vmTemplateDao.findById(anyLong())).thenReturn(mockTemplate);
        when(dataStoreManager.getPrimaryDataStore(anyLong())).thenReturn(mockPrimaryDataStore);
        when(vmTemplateDao.findById(anyLong(), anyBoolean())).thenReturn(mockTemplate);
        when(vmTemplatePoolDao.findByPoolTemplate(anyLong(), anyLong(), nullable(String.class))).thenReturn(mockTemplateStore);
        when(primaryDataStoreDao.listByStatus(StoragePoolStatus.Up)).thenReturn(pools);

        doNothing().when(mockTemplateStore).setMarkedForGC(anyBoolean());

        ExecutorService preloadExecutor = new CustomThreadPoolExecutor(8, 8, 0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue(),
                new NamedThreadFactory("Template-Preloader"));
        templateManager._preloadExecutor = preloadExecutor;

        templateManager.prepareTemplate(202, 1, null);
        assertTrue("Test template is scheduled for seeding to on pool", ((CustomThreadPoolExecutor) preloadExecutor).getCount() == 2);
    }

    @Test
    public void testCreatePrivateTemplateRecordForRegionStore() throws ResourceAllocationException {

        CreateTemplateCmd mockCreateCmd = mock(CreateTemplateCmd.class);
        when(mockCreateCmd.getTemplateName()).thenReturn("test");
        when(mockCreateCmd.getTemplateTag()).thenReturn(null);
        when(mockCreateCmd.getBits()).thenReturn(64);
        when(mockCreateCmd.getRequiresHvm()).thenReturn(true);
        when(mockCreateCmd.isPasswordEnabled()).thenReturn(false);
        when(mockCreateCmd.isPublic()).thenReturn(false);
        when(mockCreateCmd.isFeatured()).thenReturn(false);
        when(mockCreateCmd.isDynamicallyScalable()).thenReturn(false);
        when(mockCreateCmd.getVolumeId()).thenReturn(null);
        when(mockCreateCmd.getSnapshotId()).thenReturn(1L);
        when(mockCreateCmd.getOsTypeId()).thenReturn(1L);
        when(mockCreateCmd.getDetails()).thenReturn(null);
        when(mockCreateCmd.getZoneId()).thenReturn(null);

        Account mockTemplateOwner = mock(Account.class);

        SnapshotVO mockSnapshot = mock(SnapshotVO.class);
        when(snapshotDao.findById(anyLong())).thenReturn(mockSnapshot);

        when(mockSnapshot.getVolumeId()).thenReturn(1L);
        when(mockSnapshot.getState()).thenReturn(Snapshot.State.BackedUp);
        when(mockSnapshot.getHypervisorType()).thenReturn(Hypervisor.HypervisorType.XenServer);

        GuestOSVO mockGuestOS = mock(GuestOSVO.class);
        when(guestOSDao.findById(anyLong())).thenReturn(mockGuestOS);

        when(vmTemplateDao.getNextInSequence(eq(Long.class), eq("id"))).thenReturn(1L);

        List<ImageStoreVO> mockRegionStores = new ArrayList<>();
        ImageStoreVO mockRegionStore = mock(ImageStoreVO.class);
        mockRegionStores.add(mockRegionStore);
        when(imgStoreDao.findRegionImageStores()).thenReturn(mockRegionStores);

        when(vmTemplateDao.persist(any(VMTemplateVO.class))).thenAnswer(new Answer<VMTemplateVO>() {
            @Override
            public VMTemplateVO answer(InvocationOnMock invocationOnMock) throws Throwable {
                Object[] args = invocationOnMock.getArguments();
                return (VMTemplateVO)args[0];
            }
        });

        try (MockedConstruction<CheckedReservation> mockCheckedReservation = Mockito.mockConstruction(CheckedReservation.class)) {
            VMTemplateVO template = templateManager.createPrivateTemplateRecord(mockCreateCmd, mockTemplateOwner);
            assertTrue("Template in a region store should have cross zones set", template.isCrossZones());
        }
    }

    @Test
    public void testLinkUserDataToTemplate() {
        LinkUserDataToTemplateCmd cmd = Mockito.mock(LinkUserDataToTemplateCmd.class);
        when(cmd.getTemplateId()).thenReturn(1L);
        when(cmd.getIsoId()).thenReturn(null);
        when(cmd.getUserdataId()).thenReturn(2L);
        when(cmd.getUserdataPolicy()).thenReturn(UserData.UserDataOverridePolicy.ALLOWOVERRIDE);

        VMTemplateVO template = Mockito.mock(VMTemplateVO.class);
        when(vmTemplateDao.findById(anyLong())).thenReturn(template);

        when(userDataDaoMock.findById(anyLong())).thenReturn(userDataMock);

        VirtualMachineTemplate resultTemplate = templateManager.linkUserDataToTemplate(cmd);

        Assert.assertEquals(template, resultTemplate);
    }

    @Test(expected = InvalidParameterValueException.class)
    public void testLinkUserDataToTemplateByProvidingBothISOAndTemplateId() {
        LinkUserDataToTemplateCmd cmd = Mockito.mock(LinkUserDataToTemplateCmd.class);
        when(cmd.getTemplateId()).thenReturn(1L);
        when(cmd.getIsoId()).thenReturn(1L);
        when(cmd.getUserdataId()).thenReturn(2L);
        when(cmd.getUserdataPolicy()).thenReturn(UserData.UserDataOverridePolicy.ALLOWOVERRIDE);

        VMTemplateVO template = Mockito.mock(VMTemplateVO.class);

        templateManager.linkUserDataToTemplate(cmd);
    }

    @Test(expected = InvalidParameterValueException.class)
    public void testLinkUserDataToTemplateByNotProvidingBothISOAndTemplateId() {
        LinkUserDataToTemplateCmd cmd = Mockito.mock(LinkUserDataToTemplateCmd.class);
        when(cmd.getTemplateId()).thenReturn(null);
        when(cmd.getIsoId()).thenReturn(null);
        when(cmd.getUserdataId()).thenReturn(2L);
        when(cmd.getUserdataPolicy()).thenReturn(UserData.UserDataOverridePolicy.ALLOWOVERRIDE);

        VMTemplateVO template = Mockito.mock(VMTemplateVO.class);

        templateManager.linkUserDataToTemplate(cmd);
    }

    @Test(expected = InvalidParameterValueException.class)
    public void testLinkUserDataToTemplateWhenNoTemplate() {
        LinkUserDataToTemplateCmd cmd = Mockito.mock(LinkUserDataToTemplateCmd.class);
        when(cmd.getTemplateId()).thenReturn(1L);
        when(cmd.getIsoId()).thenReturn(null);
        when(cmd.getUserdataId()).thenReturn(2L);
        when(cmd.getUserdataPolicy()).thenReturn(UserData.UserDataOverridePolicy.ALLOWOVERRIDE);

        when(vmTemplateDao.findById(anyLong())).thenReturn(null);

        templateManager.linkUserDataToTemplate(cmd);
    }

    @Test
    public void testUnLinkUserDataToTemplate() {
        LinkUserDataToTemplateCmd cmd = Mockito.mock(LinkUserDataToTemplateCmd.class);
        when(cmd.getTemplateId()).thenReturn(1L);
        when(cmd.getIsoId()).thenReturn(null);
        when(cmd.getUserdataId()).thenReturn(null);
        when(cmd.getUserdataPolicy()).thenReturn(UserData.UserDataOverridePolicy.ALLOWOVERRIDE);

        VMTemplateVO template = Mockito.mock(VMTemplateVO.class);
        when(template.getId()).thenReturn(1L);
        when(vmTemplateDao.findById(1L)).thenReturn(template);

        VirtualMachineTemplate resultTemplate = templateManager.linkUserDataToTemplate(cmd);

        Assert.assertEquals(template, resultTemplate);
    }

    @Test
    public void getImageStoreTestStoreUuidIsNotNullShouldReturnAValidImageStoreIfValidUuid() {
        DataStore dataStore = Mockito.mock(DataStore.class);
        VolumeVO volumeVO = Mockito.mock(VolumeVO.class);

        Mockito.when(dataStoreManager.getDataStore(Mockito.anyString(), Mockito.any(DataStoreRole.class))).thenReturn(dataStore);

        templateManager.getImageStore("UUID", 1L, volumeVO);
    }

    @Test(expected = CloudRuntimeException.class)
    public void getImageStoreTestStoreUuidIsNotNullShouldThrowCloudRuntimeExceptionIfInvalidUuid() {
        VolumeVO volumeVO = Mockito.mock(VolumeVO.class);

        Mockito.when(dataStoreManager.getDataStore(Mockito.anyString(), Mockito.any(DataStoreRole.class))).thenReturn(null);

        templateManager.getImageStore("UUID", 1L, volumeVO);
    }

    @Test
    public void getImageStoreTestStoreUuidIsNullAndThereIsNoActiveHeuristicRulesShouldCallGetImageStoreWithFreeCapacity() {
        DataStore dataStore = Mockito.mock(DataStore.class);
        VolumeVO volumeVO = Mockito.mock(VolumeVO.class);

        Mockito.when(heuristicRuleHelperMock.getImageStoreIfThereIsHeuristicRule(Mockito.anyLong(), Mockito.any(HeuristicType.class), Mockito.any(VolumeVO.class))).thenReturn(null);
        Mockito.when(dataStoreManager.getImageStoreWithFreeCapacity(Mockito.anyLong())).thenReturn(dataStore);

        templateManager.getImageStore(null, 1L, volumeVO);
        Mockito.verify(dataStoreManager, Mockito.times(1)).getImageStoreWithFreeCapacity(Mockito.anyLong());
    }

    @Test
    public void getImageStoreTestStoreUuidIsNullAndThereIsActiveHeuristicRulesShouldNotCallGetImageStoreWithFreeCapacity() {
        DataStore dataStore = Mockito.mock(DataStore.class);
        VolumeVO volumeVO = Mockito.mock(VolumeVO.class);

        Mockito.when(heuristicRuleHelperMock.getImageStoreIfThereIsHeuristicRule(Mockito.anyLong(), Mockito.any(HeuristicType.class), Mockito.any(VolumeVO.class))).thenReturn(dataStore);

        templateManager.getImageStore(null, 1L, volumeVO);
        Mockito.verify(dataStoreManager, Mockito.times(0)).getImageStoreWithFreeCapacity(Mockito.anyLong());
    }

    private HypervisorTemplateAdapter mockHypervisorTemplateAdapter() {
        HypervisorTemplateAdapter adapter = Mockito.mock(HypervisorTemplateAdapter.class);
        when(adapter.getName()).thenReturn(TemplateAdapter.TemplateAdapterType.Hypervisor.getName());
        templateManager.setTemplateAdapters(new ArrayList<>(List.of(adapter)));
        return adapter;
    }

    @Test
    public void testRegisterTemplateValidatesUrlBeforeProbingSize() {
        RegisterTemplateCmd cmd = Mockito.mock(RegisterTemplateCmd.class);
        when(cmd.getTemplateTag()).thenReturn(null);
        when(cmd.isRoutingType()).thenReturn(null);
        when(cmd.getUrl()).thenReturn("http://10.0.0.5/template.qcow2");
        when(cmd.getFormat()).thenReturn("QCOW2");
        when(cmd.getHypervisor()).thenReturn(Hypervisor.HypervisorType.KVM.toString());
        when(cmd.isDirectDownload()).thenReturn(false);
        when(cmd.getEntityOwnerId()).thenReturn(1L);
        when(_accountMgr.getAccount(1L)).thenReturn(Mockito.mock(Account.class));

        mockHypervisorTemplateAdapter();

        try (MockedStatic<UriUtils> uriUtilsMock = Mockito.mockStatic(UriUtils.class)) {
            uriUtilsMock.when(() -> UriUtils.validateUrl(any(), any(), Mockito.anyBoolean(), Mockito.anyBoolean()))
                    .thenThrow(new IllegalArgumentException("Illegal host specified in URL."));

            Assert.assertThrows(IllegalArgumentException.class, () -> templateManager.registerTemplate(cmd));

            uriUtilsMock.verify(() -> UriUtils.getRemoteSize(any(), any()), Mockito.never());
        }
    }

    @Test
    public void testRegisterIsoValidatesUrlBeforeProbingSize() throws Exception {
        RegisterIsoCmd cmd = Mockito.mock(RegisterIsoCmd.class);
        when(cmd.getUrl()).thenReturn("http://10.0.0.5/image.iso");
        when(cmd.isDirectDownload()).thenReturn(false);
        when(cmd.getEntityOwnerId()).thenReturn(1L);
        when(_accountMgr.getAccount(1L)).thenReturn(Mockito.mock(Account.class));

        mockHypervisorTemplateAdapter();

        try (MockedStatic<UriUtils> uriUtilsMock = Mockito.mockStatic(UriUtils.class)) {
            uriUtilsMock.when(() -> UriUtils.validateUrl(any(), any(), Mockito.anyBoolean(), Mockito.anyBoolean()))
                    .thenThrow(new IllegalArgumentException("Illegal host specified in URL."));

            Assert.assertThrows(IllegalArgumentException.class, () -> templateManager.registerIso(cmd));

            uriUtilsMock.verify(() -> UriUtils.getRemoteSize(any(), any()), Mockito.never());
        }
    }

    @Test
    public void testRegisterTemplateWithTemplateType() {
        RegisterTemplateCmd cmd = Mockito.mock(RegisterTemplateCmd.class);
        when(cmd.getTemplateType()).thenReturn(Storage.TemplateType.SYSTEM.toString());
        Storage.TemplateType type = templateManager.validateTemplateType(cmd, true, true);
        Assert.assertEquals(Storage.TemplateType.SYSTEM, type);
    }

    @Test
    public void testRegisterTemplateWithoutTemplateType() {
        RegisterTemplateCmd cmd = Mockito.mock(RegisterTemplateCmd.class);
        Storage.TemplateType type = templateManager.validateTemplateType(cmd, true, true);
        Assert.assertEquals(Storage.TemplateType.USER, type);
    }

    @Test(expected = InvalidParameterValueException.class)
    public void testRegisterTemplateWithSystemTemplateTypeByUser() {
        RegisterVnfTemplateCmd cmd = Mockito.mock(RegisterVnfTemplateCmd.class);
        when(cmd.getTemplateType()).thenReturn(Storage.TemplateType.SYSTEM.toString());
        Storage.TemplateType type = templateManager.validateTemplateType(cmd, false, true);
    }

    @Test(expected = InvalidParameterValueException.class)
    public void testRegisterVnfTemplateWithTemplateType() {
        RegisterVnfTemplateCmd cmd = Mockito.mock(RegisterVnfTemplateCmd.class);
        when(cmd.getTemplateType()).thenReturn(Storage.TemplateType.SYSTEM.toString());
        Storage.TemplateType type = templateManager.validateTemplateType(cmd, false, true);
        Assert.assertEquals(Storage.TemplateType.VNF, type);
    }

    @Test
    public void testRegisterVnfTemplateWithoutTemplateType() {
        RegisterVnfTemplateCmd cmd = Mockito.mock(RegisterVnfTemplateCmd.class);
        Storage.TemplateType type = templateManager.validateTemplateType(cmd, false, true);
        Assert.assertEquals(Storage.TemplateType.VNF, type);
    }

    @Test
    public void testUpdateTemplateWithTemplateType() {
        UpdateTemplateCmd cmd = Mockito.mock(UpdateTemplateCmd.class);
        when(cmd.getTemplateType()).thenReturn(Storage.TemplateType.SYSTEM.toString());
        Storage.TemplateType type = templateManager.validateTemplateType(cmd, true, true);
        Assert.assertEquals(Storage.TemplateType.SYSTEM, type);
    }

    @Test
    public void testUpdateTemplateWithoutTemplateType() {
        UpdateTemplateCmd cmd = Mockito.mock(UpdateTemplateCmd.class);
        Storage.TemplateType type = templateManager.validateTemplateType(cmd, true, true);
        Assert.assertNull(type);
    }

    @Test(expected = InvalidParameterValueException.class)
    public void testUpdateTemplateWithInvalidTemplateType() {
        UpdateTemplateCmd cmd = Mockito.mock(UpdateTemplateCmd.class);
        when(cmd.getTemplateType()).thenReturn("invalidtype");
        Storage.TemplateType type = templateManager.validateTemplateType(cmd, true, true);
    }

    @Test(expected = InvalidParameterValueException.class)
    public void testUpdateTemplateWithInvalidTemplateTypeForRouting() {
        UpdateTemplateCmd cmd = Mockito.mock(UpdateTemplateCmd.class);
        when(cmd.getTemplateType()).thenReturn(Storage.TemplateType.USER.toString());
        when(cmd.isRoutingType()).thenReturn(true);
        Storage.TemplateType type = templateManager.validateTemplateType(cmd, true, true);
    }

    @Test(expected = InvalidParameterValueException.class)
    public void testUpdateTemplateWithInvalidCrossZonesForSystem() {
        UpdateTemplateCmd cmd = Mockito.mock(UpdateTemplateCmd.class);
        when(cmd.getTemplateType()).thenReturn(Storage.TemplateType.SYSTEM.toString());
        Storage.TemplateType type = templateManager.validateTemplateType(cmd, true, false);
    }

    @Test(expected = InvalidParameterValueException.class)
    public void testUpdateTemplateWithSystemTemplateTypeByUser() {
        UpdateVnfTemplateCmd cmd = Mockito.mock(UpdateVnfTemplateCmd.class);
        when(cmd.getTemplateType()).thenReturn(Storage.TemplateType.SYSTEM.toString());
        Storage.TemplateType type = templateManager.validateTemplateType(cmd, false, true);
    }

    @Test(expected = InvalidParameterValueException.class)
    public void testUpdateVnfTemplateWithTemplateType() {
        UpdateVnfTemplateCmd cmd = Mockito.mock(UpdateVnfTemplateCmd.class);
        when(cmd.getTemplateType()).thenReturn(Storage.TemplateType.SYSTEM.toString());
        Storage.TemplateType type = templateManager.validateTemplateType(cmd, false, true);
        Assert.assertEquals(Storage.TemplateType.VNF, type);
    }

    @Test
    public void testUpdateVnfTemplateWithoutTemplateType() {
        UpdateVnfTemplateCmd cmd = Mockito.mock(UpdateVnfTemplateCmd.class);
        Storage.TemplateType type = templateManager.validateTemplateType(cmd, false, true);
        Assert.assertNull(type);
    }

    @Test
    public void testDeleteTemplateWithTemplateType() {
        DeleteTemplateCmd cmd = new DeleteTemplateCmd();
        Storage.TemplateType type = templateManager.validateTemplateType(cmd, true, true);
        Assert.assertNull(type);
    }

    @Test
    public void verifyHeuristicRulesForZoneTestTemplateIsISOFormatShouldCheckForISOHeuristicType() {
        VMTemplateVO vmTemplateVOMock = Mockito.mock(VMTemplateVO.class);

        Mockito.when(vmTemplateVOMock.getFormat()).thenReturn(Storage.ImageFormat.ISO);
        templateManager.verifyHeuristicRulesForZone(vmTemplateVOMock, 1L);

        Mockito.verify(heuristicRuleHelperMock, Mockito.times(1)).getImageStoreIfThereIsHeuristicRule(1L, HeuristicType.ISO, vmTemplateVOMock);
    }

    @Test
    public void verifyHeuristicRulesForZoneTestTemplateNotISOFormatShouldCheckForTemplateHeuristicType() {
        VMTemplateVO vmTemplateVOMock = Mockito.mock(VMTemplateVO.class);

        Mockito.when(vmTemplateVOMock.getFormat()).thenReturn(Storage.ImageFormat.QCOW2);
        templateManager.verifyHeuristicRulesForZone(vmTemplateVOMock, 1L);

        Mockito.verify(heuristicRuleHelperMock, Mockito.times(1)).getImageStoreIfThereIsHeuristicRule(1L, HeuristicType.TEMPLATE, vmTemplateVOMock);
    }

}