package com.atlassian.vcache.internal.core.service;
   
   import com.atlassian.vcache.RequestCache;
   import com.atlassian.vcache.internal.RequestContext;
   import com.atlassian.vcache.internal.core.WorkContextRequestContextSupplier;
   import com.atlassian.workcontext.api.EstablishWorkContext;
   import com.atlassian.workcontext.api.WorkContextFuture;
   import org.junit.Rule;
   import org.junit.Test;
  import org.junit.rules.ExpectedException;
  
  import java.util.Optional;
  
  import static org.hamcrest.Matchers.is;
  import static org.hamcrest.Matchers.not;
  import static org.hamcrest.Matchers.notNullValue;
  import static org.junit.Assert.assertThat;
  
  public class WorkContextRequestContextSupplierTest {
      @Rule
      public ExpectedException thrown = ExpectedException.none();
  
      @Rule
      public EstablishWorkContext establishWorkContext = new EstablishWorkContext();
  
      @Test
      public void normal_flow() {
          final WorkContextRequestContextSupplier supplier = new WorkContextRequestContextSupplier(() -> "abc");
  
          final RequestContext context = supplier.get();
  
          assertThat(context, notNullValue());
          assertThat(context.partitionIdentifier(), is("abc"));
      }
  
      @Test
      public void add_to_cache_across_threads() {
          final RequestCache<String, String> cache = new DefaultRequestCache<>("cache", new WorkContextRequestContextSupplier(() -> "abc"));
          cache.put("1", "2");
          final String threadName = Thread.currentThread().getName();
  
          // Test in another thread.
          WorkContextFuture.runAsync( () -> {
              assertThat(Thread.currentThread().getName(), not(threadName));  // Definitely a new thread
              assertThat(cache.get("1"), is(Optional.of("2")));
              cache.put("2", "2");
          }).join();
  
          assertThat(cache.get("2"), is(Optional.of("2")));
      }
  }