Add missing headers for TB2 tests

Change-Id: Ifb5ed47f7d16630c55b5e6b5d8bc6e174c39009c