TaskContextImpl implements TaskContext { @Override public KeyValueStore getStore(String storeName) { KeyValueStore store = this.keyValueStoreProvider.apply(storeName); if (store == null) { throw new IllegalArgumentException(String.format("No store found for storeName: %s", storeName)); } return store; } TaskContextImpl(TaskModel taskModel,
      MetricsRegistry taskMetricsRegistry,
      Function<String, KeyValueStore> keyValueStoreProvider,
      TableManager tableManager,
      CallbackScheduler callbackScheduler,
      OffsetManager offsetManager,
      JobModel jobModel,
      StreamMetadataCache streamMetadataCache); @Override TaskModel getTaskModel(); @Override MetricsRegistry getTaskMetricsRegistry(); @Override KeyValueStore getStore(String storeName); @Override ReadWriteTable<K, V> getTable(String tableId); @Override CallbackScheduler getCallbackScheduler(); @Override void setStartingOffset(SystemStreamPartition systemStreamPartition, String offset); JobModel getJobModel(); StreamMetadataCache getStreamMetadataCache();  }
@Test public void testGetStore() { KeyValueStore store = mock(KeyValueStore.class); when(keyValueStoreProvider.apply("myStore")).thenReturn(store); assertEquals(store, taskContext.getStore("myStore")); }
@Test(expected = IllegalArgumentException.class) public void testGetMissingStore() { KeyValueStore store = mock(KeyValueStore.class); when(keyValueStoreProvider.apply("myStore")).thenReturn(null); assertEquals(store, taskContext.getStore("myStore")); }