From 46ded2864dc77ffed37aa9c9ff5e6dfad5883bf2 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Tue, 28 Mar 2023 14:33:33 +0300
Subject: [PATCH] Fix restart feature should be tested against agent session

---
 clearml_agent/commands/worker.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/clearml_agent/commands/worker.py b/clearml_agent/commands/worker.py
index 3c83839..b35632f 100644
--- a/clearml_agent/commands/worker.py
+++ b/clearml_agent/commands/worker.py
@@ -736,9 +736,9 @@ class Worker(ServiceCommandSection):
         except Exception:
             pass
 
-    def _get_docker_restart_value(self, session, task_id: str):
+    def _get_docker_restart_value(self, task_session, task_id: str):
         try:
-            session.verify_feature_set('advanced')
+            self._session.verify_feature_set('advanced')
         except ValueError:
             return
 
@@ -760,7 +760,7 @@ class Worker(ServiceCommandSection):
 
         path, _, expected_value = selector.partition("=")
 
-        result = session.send_request(
+        result = task_session.send_request(
             service='tasks',
             action='get_all',
             json={'id': [task_id], 'only_fields': [path], 'search_hidden': True},