Demo code cleanup