Nachdem Klein fünf Länderspiele für die U-19- und 23 für die U-21-Nationalmannschaft (drei Tore) absolviert hatte, gab er am 19. Mai 2010 sein Debüt in der A-Nationalmannschaft, die in Klagenfurt das Test-Länderspiel gegen die kroatische Auswahl mit 0:1 verlor. Im Herbst 2015 qualifizierte er sich mit der Nationalmannschaft für die Teilnahme an der EM-Endrunde 2016 in Frankreich; die erste Teilnahme einer österreichischen Nationalmannschaft an einer EM-Endrunde nach sportlicher Qualifikation. Er wurde in das EM-Aufgebot Österreichs aufgenommen und stand in allen drei Partien der Gruppenphase in der Stammelf. Er gehörte zu den Spielern, die alle Begegnungen über die volle Spielzeit absolvierten. Danach schied das Team als Gruppenletzter aus.