@Test
public void testModelIsTrained() {

    // Given
    FederatedModel model = mock(FederatedModel.class);
    FederatedDataSet dataSource = mock(FederatedDataSet.class);
    // When
    TrainModel cut = new TrainModel(model, dataSource, useCaseCallback);
    cut.execute();
    // Then
    verify(model).buildModel();
    verify(model).train(dataSource);
    verify(useCaseCallback).onSuccess(Boolean.TRUE);
}