diff --git a/extended/code_update/code_update.robot b/extended/code_update/code_update.robot
index 0d57788..73e86b6 100644
--- a/extended/code_update/code_update.robot
+++ b/extended/code_update/code_update.robot
@@ -14,9 +14,7 @@
 ...                 - Wait for code update to complete
 ...                 - Verify the new version
 
-#TODO: Move test_uploadimage.py to lib/
-Library           ../test_uploadimage.py
-Library           code_update.py
+Library           ../../lib/code_update_utils.py
 Library           OperatingSystem
 Variables         ../../data/variables.py
 Resource          ../lib/rest_client.robot