@Test
public void testGetTrainingData() {

    // Given
    FederatedDataSource dataSource = Mockito.mock(FederatedDataSource.class);
    FederatedRepository repository = Mockito.mock(FederatedRepository.class);
    given(dataSource.getTrainingData()).willReturn(dataSet);
    given(repository.getTrainingData()).willReturn(dataSet);
    // When
    GetTrainingData cut = new GetTrainingData(useCaseCallback, repository);
    cut.execute();
    // Then
    Mockito.verify(repository).getTrainingData();
    Mockito.verify(useCaseCallback).onSuccess(repository);
}