    @Test
    public void when_callingApplyGradient_expect_correctGradientAppliedAndIterationUpdated() {

        // Arrange
        setup(false);
        Gradients gradientsMock = mock(Gradients.class);
        final Gradient gradient = mock(Gradient.class);
        INDArray gradientGradient = Nd4j.rand(1, 2);
        when(gradient.gradient()).thenReturn(gradientGradient);
        when(gradientsMock.getGradient(GRADIENT_NAME)).thenReturn(gradient);
        Updater updaterMock = mock(Updater.class);
        when(modelMock.getUpdater()).thenReturn(updaterMock);
        INDArray paramsMock = mock(INDArray.class);
        when(modelMock.params()).thenReturn(paramsMock);

        // Act
        sut.applyGradient(gradientsMock, 345);

        // Assert
        verify(gradientsMock, times(1)).getGradient(GRADIENT_NAME);
        verify(updaterMock, times(1)).update(eq(modelMock), eq(gradient), eq(123), eq(234), eq(345), any());
        verify(paramsMock, times(1)).subi(gradientGradient);
        verify(configurationMock, times(1)).setIterationCount(124);
    }