How do I make sure my cosine-weighted sampling works fine?
The code you've shown here looks right to me. In order to verify the implementations, you could try writing some code that just integrates a few test functions over the hemisphere using both methods and verify that you get the same result (within some sampling error). Also if you haven't already, verify that the outputs from sample_unit_square and sample_unit_disc are doing what they're supposed to.