private void setUpTreeGrid() {
    treeGrid = new TreeGrid<>();
    treeGrid.setSizeFull();

    treeGrid.addColumn(EmployeeDto::getEmployee).setCaption("Employee").setId("employee-column");
    treeGrid.addColumn(EmployeeDto::getComputer).setCaption("Computers").setId("computer-column");
    treeGrid.addComponentColumn(this::deleteButton).setCaption("Delete"); // Generates a button on parent and child row

    data = new TreeData<>();
    data.addItems(generateEmployee(), EmployeeDto::getSubEmployee);
    dataProvider = new TreeDataProvider<>(data);
    treeGrid.setDataProvider(dataProvider);


    addComponent(treeGrid);
}

private Button deleteButton(EmployeeDto employeeDto) {

    Button button = new Button(VaadinIcons.CLOSE);
    button.addStyleName(ValoTheme.BUTTON_SMALL);
    button.addClickListener(e -> {
     //   Delete function
        onChange();
    });
    return button;
}

private List<EmployeeDto> generateEmployee() {
    List<Employee> employees = employeeService.findAllEmployee();

    List<EmployeeDto> employeeList = new ArrayList<>();
    for (Employee employee : employees) {
        EmployeeDto employeeDto;
        employeeDto = new EmployeeDto(employee.getUserName() + " " + employee.getFirstName() + " " + employee.getLastName(), String.valueOf(employee.getComputers().size()));
        if (!(employee.getComputers().isEmpty())) {
            employeeList.add(employeeDto);
        }
        for (int i = 0; i < employee.getComputers().size(); i++) {
            EmployeeDto subEmployee = new EmployeeDto();
            subEmployee.setComputer(employee.getComputers().get(i).getInventoryId());
            employeeDto.addSubEmployee(subEmployee);

        }
    }
    return employeeList;
}