Fix python batching test

Review 1679 by @larkost