/*-
 * -\-\-
 * Spotify Styx Service Common
 * --
 * Copyright (C) 2016 - 2019 Spotify AB
 * --
 * 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.spotify.styx.state;

import static com.spotify.styx.testdata.TestData.WORKFLOW_INSTANCE;
import static org.hamcrest.Matchers.is;
import static org.junit.Assert.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.when;

import com.spotify.styx.state.RunState.State;
import java.util.concurrent.atomic.AtomicReference;
import junitparams.JUnitParamsRunner;
import junitparams.Parameters;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.slf4j.MDC;

@RunWith(JUnitParamsRunner.class)
public class MDCDecoratingHandlerTest {

  private static final long COUNTER = 17L;

  @Mock private OutputHandler delegate;
  @Mock private RunState runState;

  private MDCDecoratingHandler sut;

  @Before
  public void setUp() throws Exception {
    MockitoAnnotations.initMocks(this);
    when(runState.workflowInstance()).thenReturn(WORKFLOW_INSTANCE);
    sut = new MDCDecoratingHandler(delegate);
  }

  @Parameters(source = State.class)
  @Test
  public void shouldDecorateMDC(State state) {
    when(runState.state()).thenReturn(state);
    when(runState.counter()).thenReturn(COUNTER);

    var id = new AtomicReference<String>();
    var stateCounter = new AtomicReference<String>();
    var stateName = new AtomicReference<String>();

    doAnswer(a -> {
      id.set(MDC.get("wfi-id"));
      stateCounter.set(MDC.get("wfi-state-counter"));
      stateName.set(MDC.get("wfi-state-name"));
      return null;
    }).when(delegate).transitionInto(any());

    sut.transitionInto(runState);

    assertThat(id.get(), is(WORKFLOW_INSTANCE.toString()));
    assertThat(stateCounter.get(), is(String.valueOf(COUNTER)));
    assertThat(stateName.get(), is(state.toString()));
  }
}
@@ -1024,16 +1024,44 @@
 .State;%0A
+import java.io.IOException;%0A
 import j
@@ -1215,20 +1215,85 @@
 g.junit.
-Test
+Rule;%0Aimport org.junit.Test;%0Aimport org.junit.rules.ExpectedException
 ;%0Aimport
@@ -1525,16 +1525,88 @@
 = 17L;%0A%0A
+  @Rule public ExpectedException exception = ExpectedException.none();%0A%0A
   @Mock 
@@ -2697,10 +2697,540 @@
 )));%0A  %7D
+%0A%0A  @Test%0A  public void propagatesIOException() %7B%0A    when(runState.state()).thenReturn(State.RUNNING);%0A    when(runState.counter()).thenReturn(COUNTER);%0A%0A    var cause = new IOException(%22foo!%22);%0A%0A    doAnswer(a -%3E sneakyThrow(cause)).when(delegate).transitionInto(any());%0A%0A    exception.expect(RuntimeException.class);%0A    exception.expectCause(is(cause));%0A    sut.transitionInto(runState);%0A  %7D%0A%0A  @SuppressWarnings(%22unchecked%22)%0A  private static %3CE extends Throwable%3E Void sneakyThrow(Throwable e) throws E %7B%0A    throw (E) e;%0A  %7D
 %0A%7D

add test