Use pytest for testing