Add diamond test code