webdriver - how to capture screenshot on the assert statement passed

Question

@Override
public void onTestSuccess(ITestResult tr) {


}

I add method onTestSuccess, but it takes screenshot only when the whole test passed, not assetion passed.

what can I do?