Add main test function