DefaultTaskToPodConverter implements TaskToPodConverter { @VisibleForTesting void appendS3WriterRole(TitanProtos.ContainerInfo.Builder containerInfoBuilder, Job<?> job, Task task) { if (!configuration.isDefaultS3WriterRoleEnabled()) { return; } if (LogStorageInfos.findCustomS3Bucket(job).isPresent()) { containerInfoBuilder.putPassthroughAttributes( S3_WRITER_ROLE, job.getJobDescriptor().getContainer().getSecurityProfile().getIamRole() ); } else { Evaluators.applyNotNull( configuration.getDefaultS3WriterRole(), role -> containerInfoBuilder.putPassthroughAttributes(S3_WRITER_ROLE, role) ); } logStorageInfo.getS3LogLocation(task, false).ifPresent(s3LogLocation -> Evaluators.applyNotNull( s3LogLocation.getBucket(), bucket -> containerInfoBuilder.putPassthroughAttributes(S3_BUCKET_NAME, bucket) ) ); } @Inject  DefaultTaskToPodConverter(DirectKubeConfiguration configuration,
                                     PodAffinityFactory podAffinityFactory,
                                     TaintTolerationFactory taintTolerationFactory,
                                     LogStorageInfo<Task> logStorageInfo); @Override V1Pod apply(Job<?> job, Task task);  }
@Test public void testDefaultWriterRoleAssignment() { BatchJobTask task = JobGenerator.oneBatchTask(); when(logStorageInfo.getS3LogLocation(task, false)).thenReturn(Optional.of(DEFAULT_S3_LOG_LOCATION)); when(configuration.getDefaultS3WriterRole()).thenReturn("defaultWriter"); when(configuration.isDefaultS3WriterRoleEnabled()).thenReturn(true); ContainerInfo.Builder containerInfoBuilder = ContainerInfo.newBuilder(); converter.appendS3WriterRole(containerInfoBuilder, JobGenerator.oneBatchJob(), task); assertThat(containerInfoBuilder.getPassthroughAttributesMap()).containsEntry(S3_WRITER_ROLE, "defaultWriter"); assertThat(containerInfoBuilder.getPassthroughAttributesMap()).containsEntry(S3_BUCKET_NAME, "defaultBucket"); }
@Test public void testNullDefaultWriterRoleAssignment() { ContainerInfo.Builder containerInfoBuilder = ContainerInfo.newBuilder(); converter.appendS3WriterRole(containerInfoBuilder, JobGenerator.oneBatchJob(), JobGenerator.oneBatchTask()); assertThat(containerInfoBuilder.getPassthroughAttributesMap()).doesNotContainKey(S3_WRITER_ROLE); }
@Test public void testCustomBucketWriterRoleAssignment() { when(configuration.isDefaultS3WriterRoleEnabled()).thenReturn(true); ContainerInfo.Builder containerInfoBuilder = ContainerInfo.newBuilder(); Job<BatchJobExt> job = JobGenerator.oneBatchJob(); job = job.toBuilder() .withJobDescriptor(job.getJobDescriptor().toBuilder() .withContainer(job.getJobDescriptor().getContainer().toBuilder() .withAttributes(Collections.singletonMap( JobAttributes.JOB_CONTAINER_ATTRIBUTE_S3_BUCKET_NAME, "myOwnBucket" )).build() ) .build() ) .build(); converter.appendS3WriterRole(containerInfoBuilder, job, JobGenerator.oneBatchTask()); assertThat(containerInfoBuilder.getPassthroughAttributesMap()).containsEntry( S3_WRITER_ROLE, job.getJobDescriptor().getContainer().getSecurityProfile().getIamRole() ); }