Update UserVmManagerImplTest.java

gzhao9 · Oct 18, 2023 · 8530432 · 8530432
1 parent ee25fac
commit 8530432
Showing 1 changed file with 10 additions and 10 deletions.
diff --git a/server/src/test/java/com/cloud/vm/UserVmManagerImplTest.java b/server/src/test/java/com/cloud/vm/UserVmManagerImplTest.java
@@ -706,8 +706,7 @@ private ServiceOfferingVO prepareOfferingsForEncryptionValidation(long diskOffer
     public void testUserDataDenyOverride() {
         Long userDataId = 1L;
 
-        VirtualMachineTemplate template = Mockito.mock(VirtualMachineTemplate.class);
-        when(template.getUserDataId()).thenReturn(2L);
+        VirtualMachineTemplate template = CreateMockVirtualMachineTemplate(2L);
         when(template.getUserDataOverridePolicy()).thenReturn(UserData.UserDataOverridePolicy.DENYOVERRIDE);
 
         userVmManagerImpl.finalizeUserData(null, userDataId, template);
@@ -718,8 +717,7 @@ public void testUserDataAllowOverride() {
         String templateUserData = "testTemplateUserdata";
         Long userDataId = 1L;
 
-        VirtualMachineTemplate template = Mockito.mock(VirtualMachineTemplate.class);
-        when(template.getUserDataId()).thenReturn(2L);
+        VirtualMachineTemplate template = CreateMockVirtualMachineTemplate(2L);
         when(template.getUserDataOverridePolicy()).thenReturn(UserData.UserDataOverridePolicy.ALLOWOVERRIDE);
 
         UserDataVO apiUserDataVO = Mockito.mock(UserDataVO.class);
@@ -740,8 +738,7 @@ public void testUserDataWithoutTemplate() {
         doReturn(apiUserDataVO).when(userDataDao).findById(userDataId);
         when(apiUserDataVO.getUserData()).thenReturn(userData);
 
-        VirtualMachineTemplate template = Mockito.mock(VirtualMachineTemplate.class);
-        when(template.getUserDataId()).thenReturn(null);
+        VirtualMachineTemplate template = CreateMockVirtualMachineTemplate(null);
 
         String finalUserdata = userVmManagerImpl.finalizeUserData(null, userDataId, template);
 
@@ -752,8 +749,7 @@ public void testUserDataWithoutTemplate() {
     public void testUserDataAllowOverrideWithoutAPIuserdata() {
         String templateUserData = "testTemplateUserdata";
 
-        VirtualMachineTemplate template = Mockito.mock(VirtualMachineTemplate.class);
-        when(template.getUserDataId()).thenReturn(2L);
+        VirtualMachineTemplate template = CreateMockVirtualMachineTemplate(2L);
         when(template.getUserDataOverridePolicy()).thenReturn(UserData.UserDataOverridePolicy.ALLOWOVERRIDE);
         UserDataVO templateUserDataVO = Mockito.mock(UserDataVO.class);
         doReturn(templateUserDataVO).when(userDataDao).findById(2L);
@@ -767,8 +763,7 @@ public void testUserDataAllowOverrideWithoutAPIuserdata() {
     @Test
     public void testUserDataAllowOverrideWithUserdataText() {
         String userData = "testUserdata";
-        VirtualMachineTemplate template = Mockito.mock(VirtualMachineTemplate.class);
-        when(template.getUserDataId()).thenReturn(null);
+        VirtualMachineTemplate template = CreateMockVirtualMachineTemplate(null);
 
         String finalUserdata = userVmManagerImpl.finalizeUserData(userData, null, template);
 
@@ -1068,4 +1063,9 @@ public void testUpdateVncPasswordIfItHasChangedNewPassword() {
         Mockito.verify(userVmDao).findById(vmId);
         Mockito.verify(userVmDao).update(vmId, userVmVoMock);
     }
+    VirtualMachineTemplate CreateMockVirtualMachineTemplate(Long returnResult){
+        VirtualMachineTemplate template = Mockito.mock(VirtualMachineTemplate.class);
+        when(template.getUserDataId()).thenReturn(returnResult);
+        return  template;
+    }
 }