Source:
<code>public void testSpanTraceCompletionTime() throws RetryAttemptException { SpanCache spanCacheMock = Mockito.mock(SpanCache.class); CompletionTimeProcessingDeriver completionTimeProcessingDeriver = new CompletionTimeProcessingDeriver(spanCacheMock); Span rootSpan = new Span(); rootSpan.setId("trace"); rootSpan.setTraceId("trace"); rootSpan.setTimestamp(0); Annotation annotation = new Annotation(); annotation.setTimestamp(0); Annotation annotation1 = new Annotation(); annotation1.setTimestamp(1000); rootSpan.setDuration(annotation1.getTimestamp() - annotation.getTimestamp()); rootSpan.setAnnotations(Arrays.asList(annotation, annotation1)); Mockito.when(spanCacheMock.getTrace(null, "trace")).thenReturn(Arrays.asList(rootSpan)); CompletionTimeProcessing completionTimeProcessing = new CompletionTimeProcessing(rootSpan); completionTimeProcessing = completionTimeProcessingDeriver.processOneToOne(null, completionTimeProcessing); Assert.assertNull(completionTimeProcessing.getCompletionTime()); completionTimeProcessing = completionTimeProcessingDeriver.processOneToOne(null, completionTimeProcessing); Assert.assertEquals(1, completionTimeProcessing.getCompletionTime().getDuration()); <START> Span descendant = new Span(); <END> descendant.setId("descendant"); descendant.setTraceId("trace"); descendant.setTimestamp(1500); annotation = new Annotation(); annotation.setTimestamp(1500); annotation1 = new Annotation(); annotation1.setTimestamp(2000); descendant.setAnnotations(Arrays.asList(annotation, annotation1)); Mockito.when(spanCacheMock.getTrace(null, "trace")).thenReturn(Arrays.asList(rootSpan, descendant)); completionTimeProcessing = completionTimeProcessingDeriver.processOneToOne(null, completionTimeProcessing); Assert.assertNull(completionTimeProcessing.getCompletionTime()); completionTimeProcessing = completionTimeProcessingDeriver.processOneToOne(null, completionTimeProcessing); Assert.assertEquals(2, completionTimeProcessing.getCompletionTime().getDuration()); Mockito.verify(spanCacheMock, Mockito.times(4)).getTrace(null, "trace"); }</code><technical_language>this separate test. i.e. test for single span, another for more one</technical_language>

Target:
public void testSpanTraceCompletionTime() throws RetryAttemptException { SpanCache spanCacheMock = Mockito.mock(SpanCache.class); CompletionTimeProcessingDeriver completionTimeProcessingDeriver = new CompletionTimeProcessingDeriver(spanCacheMock); Span rootSpan = new Span(); rootSpan.setId("trace"); rootSpan.setTraceId("trace"); rootSpan.setTimestamp(0); Annotation annotation = new Annotation(); annotation.setTimestamp(0); Annotation annotation1 = new Annotation(); annotation1.setTimestamp(1000); rootSpan.setDuration(annotation1.getTimestamp() - annotation.getTimestamp()); rootSpan.setAnnotations(Arrays.asList(annotation, annotation1)); Mockito.when(spanCacheMock.getTrace(null, "trace")).thenReturn(new HashSet<>(Arrays.asList(rootSpan))); CompletionTimeProcessing completionTimeProcessing = new CompletionTimeProcessing(rootSpan); completionTimeProcessing = completionTimeProcessingDeriver.processOneToOne(null, completionTimeProcessing); Assert.assertNull(completionTimeProcessing.getCompletionTime()); completionTimeProcessing = completionTimeProcessingDeriver.processOneToOne(null, completionTimeProcessing); Assert.assertEquals(1, completionTimeProcessing.getCompletionTime().getDuration()); completionTimeProcessing = completionTimeProcessingDeriver.processOneToOne(null, completionTimeProcessing); Assert.assertNull(completionTimeProcessing); Span descendant = new Span(); descendant.setId("descendant"); descendant.setTraceId("trace"); descendant.setTimestamp(1500); annotation = new Annotation(); annotation.setTimestamp(1500); annotation1 = new Annotation(); annotation1.setTimestamp(2000); descendant.setAnnotations(Arrays.asList(annotation, annotation1)); Mockito.when(spanCacheMock.getTrace(null, "trace")).thenReturn( new HashSet<>(Arrays.asList(rootSpan, descendant))); completionTimeProcessing = new CompletionTimeProcessing(rootSpan); completionTimeProcessing = completionTimeProcessingDeriver.processOneToOne(null, completionTimeProcessing); Assert.assertNull(completionTimeProcessing.getCompletionTime()); completionTimeProcessing = completionTimeProcessingDeriver.processOneToOne(null, completionTimeProcessing); Assert.assertEquals(2, completionTimeProcessing.getCompletionTime().getDuration()); completionTimeProcessing = completionTimeProcessingDeriver.processOneToOne(null, completionTimeProcessing); Assert.assertNull(completionTimeProcessing); Mockito.verify(spanCacheMock, Mockito.times(4)).getTrace(null, "trace"); }