/** * Copyright 2017 The OpenZipkin Authors * * 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 zipkin.sparkstreaming.consumer.storage; import java.io.IOException; import java.util.Collections; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; import org.junit.Rule; import org.junit.Test; import org.junit.rules.ExpectedException; import org.mockito.stubbing.Answer; import org.slf4j.Logger; import slf4jtest.LogLevel; import slf4jtest.Settings; import slf4jtest.TestLogger; import zipkin.TestObjects; import zipkin.storage.AsyncSpanConsumer; import zipkin.storage.Callback; import zipkin.storage.InMemoryStorage; import zipkin.storage.StorageComponent; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.tuple; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; public class StorageConsumerTest { @Rule public ExpectedException thrown = ExpectedException.none(); StorageComponent storage = mock(StorageComponent.class); TestLogger logger = new Settings().enableAll().buildLogging().getLogger(""); StorageConsumer storageConsumer = new StorageConsumer() { @Override Logger log() { return logger; } @Override protected StorageComponent tryCompute() { return storage; } }; @Test public void logsWhenEmpty() { storageConsumer.accept(Collections.emptyList()); assertThat(logger.lines()) .extracting("level", "text") .containsExactly(tuple(LogLevel.DebugLevel, "Input was empty")); } @Test public void acceptsTrace() { storage = new InMemoryStorage(); StorageConsumer storageConsumer = new StorageConsumer() { @Override Logger log() { return logger; } @Override protected StorageComponent tryCompute() { return storage; } }; storageConsumer.accept(TestObjects.TRACE); assertThat(logger.lines()) .extracting("level", "text") .containsExactly(tuple(LogLevel.DebugLevel, "Wrote 3 spans")); assertThat(storage.spanStore().getRawTrace( TestObjects.TRACE.get(0).traceIdHigh, TestObjects.TRACE.get(0).traceId )).isEqualTo(TestObjects.TRACE); } @Test public void logsOnAcceptError() { IllegalStateException acceptException = new IllegalStateException("failed"); AsyncSpanConsumer consumer = mock(AsyncSpanConsumer.class); when(storage.asyncSpanConsumer()).thenReturn(consumer); doAnswer(answer(c -> { throw acceptException; })).when(consumer).accept(eq(TestObjects.TRACE), any(Callback.class)); storageConsumer.accept(TestObjects.TRACE); assertThat(logger.lines()).hasSize(1); assertThat(logger.lines().iterator().next().toString()) .startsWith("LogMessage(,WARN,Dropped 3 spans: failed"); // TODO: test for acceptException } @Test public void logsOnCallbackError() { IllegalStateException callbackException = new IllegalStateException("failed"); AsyncSpanConsumer consumer = mock(AsyncSpanConsumer.class); when(storage.asyncSpanConsumer()).thenReturn(consumer); doAnswer(answer(c -> c.onError(callbackException))) .when(consumer).accept(eq(TestObjects.TRACE), any(Callback.class)); storageConsumer.accept(TestObjects.TRACE); assertThat(logger.lines()).hasSize(1); assertThat(logger.lines().iterator().next().toString()) .startsWith("LogMessage(,WARN,Dropped 3 spans: failed"); // TODO: test for callbackException } @Test public void doesntWrapCheckedExceptionOnCallbackError() { IOException callbackException = new IOException("failed"); AsyncSpanConsumer consumer = mock(AsyncSpanConsumer.class); when(storage.asyncSpanConsumer()).thenReturn(consumer); doAnswer(answer(c -> c.onError(callbackException))) .when(consumer).accept(eq(TestObjects.TRACE), any(Callback.class)); storageConsumer.accept(TestObjects.TRACE); assertThat(logger.lines()).hasSize(1); assertThat(logger.lines().iterator().next().toString()) .startsWith("LogMessage(,WARN,Dropped 3 spans: failed"); // TODO: test for callbackException } @Test(timeout = 1000L) public void get_memoizes() throws InterruptedException { AtomicInteger provisionCount = new AtomicInteger(); StorageConsumer storageConsumer = new StorageConsumer() { @Override protected StorageComponent tryCompute() { provisionCount.incrementAndGet(); return new InMemoryStorage(); } }; int getCount = 1000; CountDownLatch latch = new CountDownLatch(getCount); ExecutorService exec = Executors.newFixedThreadPool(10); for (int i = 0; i < getCount; i++) { exec.execute(() -> { storageConsumer.get(); latch.countDown(); }); } latch.await(); exec.shutdown(); exec.awaitTermination(1, TimeUnit.SECONDS); assertThat(provisionCount.get()).isEqualTo(1); } static <T> Answer answer(Consumer<Callback<T>> onCallback) { return invocation -> { onCallback.accept((Callback) invocation.getArguments()[invocation.getArguments().length - 1]); return null; }; } }