Skip to content

Commit

Permalink
Issue #3534: Tag compute resources according to the tool's metadata -…
Browse files Browse the repository at this point in the history
… tests
  • Loading branch information
ekazachkova committed May 27, 2024
1 parent 22170a0 commit a159aaa
Show file tree
Hide file tree
Showing 6 changed files with 284 additions and 15 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,41 @@
import com.amazonaws.AmazonServiceException;
import com.amazonaws.services.ec2.AmazonEC2;
import com.amazonaws.services.ec2.AmazonEC2ClientBuilder;
import com.amazonaws.services.ec2.model.*;
import com.amazonaws.services.ec2.model.AmazonEC2Exception;
import com.amazonaws.services.ec2.model.AttachVolumeRequest;
import com.amazonaws.services.ec2.model.AvailabilityZone;
import com.amazonaws.services.ec2.model.CreateTagsRequest;
import com.amazonaws.services.ec2.model.CreateVolumeRequest;
import com.amazonaws.services.ec2.model.DeleteTagsRequest;
import com.amazonaws.services.ec2.model.DeleteVolumeRequest;
import com.amazonaws.services.ec2.model.DescribeImagesRequest;
import com.amazonaws.services.ec2.model.DescribeInstanceTypesRequest;
import com.amazonaws.services.ec2.model.DescribeInstanceTypesResult;
import com.amazonaws.services.ec2.model.DescribeInstancesRequest;
import com.amazonaws.services.ec2.model.DescribeNetworkInterfacesRequest;
import com.amazonaws.services.ec2.model.DescribeNetworkInterfacesResult;
import com.amazonaws.services.ec2.model.DescribeSpotPriceHistoryRequest;
import com.amazonaws.services.ec2.model.DescribeSpotPriceHistoryResult;
import com.amazonaws.services.ec2.model.DescribeVolumesRequest;
import com.amazonaws.services.ec2.model.EbsInstanceBlockDevice;
import com.amazonaws.services.ec2.model.EbsInstanceBlockDeviceSpecification;
import com.amazonaws.services.ec2.model.Filter;
import com.amazonaws.services.ec2.model.Instance;
import com.amazonaws.services.ec2.model.InstanceBlockDeviceMapping;
import com.amazonaws.services.ec2.model.InstanceBlockDeviceMappingSpecification;
import com.amazonaws.services.ec2.model.InstanceStateName;
import com.amazonaws.services.ec2.model.InstanceTypeInfo;
import com.amazonaws.services.ec2.model.ModifyInstanceAttributeRequest;
import com.amazonaws.services.ec2.model.NetworkInterface;
import com.amazonaws.services.ec2.model.Placement;
import com.amazonaws.services.ec2.model.Reservation;
import com.amazonaws.services.ec2.model.SpotPrice;
import com.amazonaws.services.ec2.model.StartInstancesRequest;
import com.amazonaws.services.ec2.model.StateReason;
import com.amazonaws.services.ec2.model.StopInstancesRequest;
import com.amazonaws.services.ec2.model.Tag;
import com.amazonaws.services.ec2.model.TerminateInstancesRequest;
import com.amazonaws.services.ec2.model.Volume;
import com.amazonaws.waiters.Waiter;
import com.amazonaws.waiters.WaiterParameters;
import com.epam.pipeline.common.MessageConstants;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -358,9 +358,7 @@ public Set<String> getMetadataKeys(final AclClass entityClass) {
public Map<String, String> prepareCustomInstanceTags(final PipelineRun run) {
try {
final Map<String, String> customTags = resolveCommonCustomInstanceTags(run);
final Tool tool = toolManager.loadByNameOrId(run.getDockerImage());
final MetadataEntry toolMetadata = loadMetadataItem(tool.getId(), AclClass.TOOL);
return resolveInstanceTagsFromMetadata(toolMetadata, customTags);
return resolveInstanceTagsFromMetadata(run.getDockerImage(), customTags);
} catch (Exception e) {
LOGGER.error("An error occurred during custom tags preparation for run '{}'.", run.getId(), e);
return new HashMap<>();
Expand Down Expand Up @@ -427,14 +425,8 @@ private void checkEntityCanBeModified(final Object entity) {
MessageConstants.ERROR_TOOL_SYMLINK_MODIFICATION_NOT_SUPPORTED)));
}

private Map<String, String> resolveInstanceTagsFromMetadata(final MetadataEntry metadataEntry,
private Map<String, String> resolveInstanceTagsFromMetadata(final String dockerImage,
final Map<String, String> customTags) {
final Map<String, PipeConfValue> metadataData = MapUtils.emptyIfNull(Objects.isNull(metadataEntry)
? null
: metadataEntry.getData());
if (MapUtils.isEmpty(metadataData)) {
return customTags;
}
final Set<String> instanceTagsKeys = new HashSet<>(Arrays.asList(preferenceManager.findPreference(
SystemPreferences.CLUSTER_INSTANCE_ALLOWED_CUSTOM_TAGS)
.filter(StringUtils::isNotBlank)
Expand All @@ -443,6 +435,17 @@ private Map<String, String> resolveInstanceTagsFromMetadata(final MetadataEntry
if (CollectionUtils.isEmpty(instanceTagsKeys)) {
return customTags;
}

final Tool tool = toolManager.loadByNameOrId(dockerImage);
final MetadataEntry toolMetadata = loadMetadataItem(tool.getId(), AclClass.TOOL);

final Map<String, PipeConfValue> metadataData = MapUtils.emptyIfNull(Objects.isNull(toolMetadata)
? null
: toolMetadata.getData());
if (MapUtils.isEmpty(metadataData)) {
return customTags;
}

metadataData.entrySet().stream()
.filter(entry -> instanceTagsKeys.contains(entry.getKey()))
.forEach(entry -> customTags.put(entry.getKey(), entry.getValue().getValue()));
Expand Down
4 changes: 0 additions & 4 deletions api/src/test/java/com/epam/pipeline/app/TestApplication.java
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@
import com.epam.pipeline.manager.datastorage.lifecycle.DataStorageLifecycleManager;
import com.epam.pipeline.manager.datastorage.providers.StorageEventCollector;
import com.epam.pipeline.manager.ldap.LdapManager;
import com.epam.pipeline.manager.metadata.MetadataManager;
import com.epam.pipeline.manager.notification.ContextualNotificationManager;
import com.epam.pipeline.manager.notification.ContextualNotificationRegistrationManager;
import com.epam.pipeline.manager.notification.ContextualNotificationSettingsManager;
Expand Down Expand Up @@ -188,9 +187,6 @@ public static void main(String[] args) {
@MockBean
public InstanceOfferScheduler instanceOfferScheduler;

@MockBean
public MetadataManager metadataManager;

@Bean
public EmbeddedServletContainerCustomizer containerCustomizer() throws FileNotFoundException {

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,230 @@
/*
* Copyright 2017-2024 EPAM Systems, Inc. (https://www.epam.com/)
*
* 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.
*/

package com.epam.pipeline.manager.metadata;

import com.epam.pipeline.controller.vo.EntityVO;
import com.epam.pipeline.dao.metadata.MetadataDao;
import com.epam.pipeline.entity.metadata.CommonCustomInstanceTagsTypes;
import com.epam.pipeline.entity.metadata.MetadataEntry;
import com.epam.pipeline.entity.metadata.PipeConfValue;
import com.epam.pipeline.entity.pipeline.PipelineRun;
import com.epam.pipeline.entity.pipeline.Tool;
import com.epam.pipeline.entity.security.acl.AclClass;
import com.epam.pipeline.manager.pipeline.ToolManager;
import com.epam.pipeline.manager.preference.PreferenceManager;
import com.epam.pipeline.manager.preference.SystemPreferences;
import com.google.common.collect.Maps;
import org.junit.Before;
import org.junit.Test;
import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;

import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Matchers.any;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

public class MetadataManagerUnitTest {
private static final String KEY_1 = "key1";
private static final String KEY_2 = "key2";
private static final String KEY_3 = "key3";
private static final String VALUE_1 = "value1";
private static final String TYPE = "string";
private static final String TEST_USER = "TEST_USER";
private static final String OWNER_KEY = "CP_OWNER";
private static final String RUN_ID_KEY = "CP_RUN_ID";
private static final String TOOL_KEY = "CP_TOOL";
private static final String TEST_RUN_ID = "1";
private static final String TEST_IMAGE = "test:8080/docker/image:latest";
private static final long TEST_TOOL_ID = 1L;

@Mock
private ToolManager toolManager;
@Mock
private PreferenceManager preferenceManager;
@Mock
private MetadataDao metadataDao;

@InjectMocks
private MetadataManager metadataManager;

@Before
public void setUp() throws Exception {
MockitoAnnotations.initMocks(this);
}

@Test
public void shouldPrepareTagsFromPreferenceAndToolsMetadata() {
when(preferenceManager.getPreference(SystemPreferences.CLUSTER_INSTANCE_TAGS))
.thenReturn(buildCommonTagsMapping());

final PipelineRun run = new PipelineRun();
run.setId(Long.valueOf(TEST_RUN_ID));
run.setOwner(TEST_USER);
run.setDockerImage(TEST_IMAGE);

final Tool tool = new Tool();
tool.setId(TEST_TOOL_ID);
when(toolManager.loadByNameOrId(TEST_IMAGE)).thenReturn(tool);

when(preferenceManager.findPreference(SystemPreferences.CLUSTER_INSTANCE_ALLOWED_CUSTOM_TAGS))
.thenReturn(Optional.of(String.join(",", KEY_1, KEY_2)));

final EntityVO entityVO = new EntityVO(TEST_TOOL_ID, AclClass.TOOL);
when(metadataDao.loadMetadataItem(entityVO)).thenReturn(toolMetadata(entityVO));

final Map<String, String> tags = metadataManager.prepareCustomInstanceTags(run);
assertThat(tags)
.hasSize(4)
.contains(Maps.immutableEntry(KEY_1, VALUE_1))
.contains(Maps.immutableEntry(OWNER_KEY, TEST_USER))
.contains(Maps.immutableEntry(RUN_ID_KEY, TEST_RUN_ID))
.contains(Maps.immutableEntry(TOOL_KEY, TEST_IMAGE));
}

@Test
public void shouldPrepareTagsFromPreference() {
when(preferenceManager.getPreference(SystemPreferences.CLUSTER_INSTANCE_TAGS))
.thenReturn(buildCommonTagsMapping());

final PipelineRun run = new PipelineRun();
run.setId(Long.valueOf(TEST_RUN_ID));
run.setOwner(TEST_USER);
run.setDockerImage(TEST_IMAGE);

when(preferenceManager.findPreference(SystemPreferences.CLUSTER_INSTANCE_ALLOWED_CUSTOM_TAGS))
.thenReturn(Optional.empty());

final Map<String, String> tags = metadataManager.prepareCustomInstanceTags(run);
assertThat(tags)
.hasSize(3)
.contains(Maps.immutableEntry(OWNER_KEY, TEST_USER))
.contains(Maps.immutableEntry(RUN_ID_KEY, TEST_RUN_ID))
.contains(Maps.immutableEntry(TOOL_KEY, TEST_IMAGE));
verify(metadataDao, never()).loadMetadataItem(any());
verify(toolManager, never()).loadByNameOrId(any());
}

@Test
public void shouldPrepareTagsFromToolsMetadata() {
when(preferenceManager.getPreference(SystemPreferences.CLUSTER_INSTANCE_TAGS)).thenReturn(null);

final PipelineRun run = new PipelineRun();
run.setId(Long.valueOf(TEST_RUN_ID));
run.setOwner(TEST_USER);
run.setDockerImage(TEST_IMAGE);

final Tool tool = new Tool();
tool.setId(TEST_TOOL_ID);
when(toolManager.loadByNameOrId(TEST_IMAGE)).thenReturn(tool);

final EntityVO entityVO = new EntityVO(TEST_TOOL_ID, AclClass.TOOL);
when(metadataDao.loadMetadataItem(entityVO)).thenReturn(toolMetadata(entityVO));

when(preferenceManager.findPreference(SystemPreferences.CLUSTER_INSTANCE_ALLOWED_CUSTOM_TAGS))
.thenReturn(Optional.of(String.join(",", KEY_1, KEY_2)));

final Map<String, String> tags = metadataManager.prepareCustomInstanceTags(run);
assertThat(tags)
.hasSize(1)
.contains(Maps.immutableEntry(KEY_1, VALUE_1));
}

@Test
public void shouldReturnEmptyTagsIfToolsMetadataNotMatch() {
when(preferenceManager.getPreference(SystemPreferences.CLUSTER_INSTANCE_TAGS)).thenReturn(null);

final PipelineRun run = new PipelineRun();
run.setId(Long.valueOf(TEST_RUN_ID));
run.setOwner(TEST_USER);
run.setDockerImage(TEST_IMAGE);

final Tool tool = new Tool();
tool.setId(TEST_TOOL_ID);
when(toolManager.loadByNameOrId(TEST_IMAGE)).thenReturn(tool);

final EntityVO entityVO = new EntityVO(TEST_TOOL_ID, AclClass.TOOL);
when(metadataDao.loadMetadataItem(entityVO)).thenReturn(toolMetadata(entityVO));

when(preferenceManager.findPreference(SystemPreferences.CLUSTER_INSTANCE_ALLOWED_CUSTOM_TAGS))
.thenReturn(Optional.of(KEY_2));

final Map<String, String> tags = metadataManager.prepareCustomInstanceTags(run);
assertThat(tags).hasSize(0);
}

@Test
public void shouldReturnEmptyTagsIfError() {
when(preferenceManager.getPreference(SystemPreferences.CLUSTER_INSTANCE_TAGS))
.thenReturn(buildCommonTagsMapping());

final PipelineRun run = new PipelineRun();
run.setId(Long.valueOf(TEST_RUN_ID));

when(preferenceManager.findPreference(SystemPreferences.CLUSTER_INSTANCE_ALLOWED_CUSTOM_TAGS))
.thenReturn(Optional.of(String.join(",", KEY_1, KEY_2)));

final Map<String, String> tags = metadataManager.prepareCustomInstanceTags(run);
assertThat(tags).hasSize(0);
}

@Test
public void shouldReturnEmptyTagsIfNoToolMetadata() {
when(preferenceManager.getPreference(SystemPreferences.CLUSTER_INSTANCE_TAGS)).thenReturn(null);

final PipelineRun run = new PipelineRun();
run.setId(Long.valueOf(TEST_RUN_ID));
run.setDockerImage(TEST_IMAGE);

final Tool tool = new Tool();
tool.setId(TEST_TOOL_ID);
when(toolManager.loadByNameOrId(TEST_IMAGE)).thenReturn(tool);

when(preferenceManager.findPreference(SystemPreferences.CLUSTER_INSTANCE_ALLOWED_CUSTOM_TAGS))
.thenReturn(Optional.of(String.join(",", KEY_1, KEY_2)));
final EntityVO entityVO = new EntityVO(TEST_TOOL_ID, AclClass.TOOL);
when(metadataDao.loadMetadataItem(entityVO)).thenReturn(null);

final Map<String, String> tags = metadataManager.prepareCustomInstanceTags(run);
assertThat(tags).hasSize(0);
verify(toolManager).loadByNameOrId(TEST_IMAGE);
}

private static Map<CommonCustomInstanceTagsTypes, String> buildCommonTagsMapping() {
final Map<CommonCustomInstanceTagsTypes, String> mapping = new HashMap<>();
mapping.put(CommonCustomInstanceTagsTypes.owner, OWNER_KEY);
mapping.put(CommonCustomInstanceTagsTypes.run_id, RUN_ID_KEY);
mapping.put(CommonCustomInstanceTagsTypes.tool, TOOL_KEY);
return mapping;
}

private static MetadataEntry toolMetadata(final EntityVO entityVO) {
final Map<String, PipeConfValue> data = new HashMap<>();
data.put(KEY_1, new PipeConfValue(TYPE, VALUE_1));
data.put(KEY_3, new PipeConfValue(TYPE, VALUE_1));
final MetadataEntry metadataEntry = new MetadataEntry();
metadataEntry.setEntity(entityVO);
metadataEntry.setData(data);
return metadataEntry;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

@SuppressWarnings("unused")
public class PipelineRunManagerUnitTest {
private static final Long RUN_ID = 1L;
private static final long NOT_EXISTING_RUN_ID = -1L;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
import com.epam.pipeline.manager.datastorage.DataStorageValidator;
import com.epam.pipeline.manager.datastorage.StorageProviderManager;
import com.epam.pipeline.manager.datastorage.providers.aws.s3.S3StorageProvider;
import com.epam.pipeline.manager.metadata.MetadataManager;
import com.epam.pipeline.manager.pipeline.FolderManager;
import com.epam.pipeline.manager.preference.PreferenceManager;
import com.epam.pipeline.manager.preference.SystemPreferences;
Expand Down Expand Up @@ -76,6 +77,7 @@
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;

@SuppressWarnings("unused")
@Transactional
public class UserManagerTest extends AbstractSpringTest {

Expand Down Expand Up @@ -129,6 +131,9 @@ public class UserManagerTest extends AbstractSpringTest {
@MockBean
private DataStorageValidator storageValidator;

@MockBean
private MetadataManager metadataManager;

@Before
public void setUpPreferenceManager() {
doReturn(mock(S3StorageProvider.class)).when(storageProviderManager).getStorageProvider(any());
Expand Down

0 comments on commit a159aaa

Please sign in to comment.