dabla commented on PR #56457:
URL: https://github.com/apache/airflow/pull/56457#issuecomment-3645409853
I've just checked the code and I saw maybe a possible "improvement"
regarding following code in worker:
```
if worker_info.state == EdgeWorkerState.MAINTENANCE_REQUEST:
logger.info("Maintenance mode requested!")
EdgeWorker.maintenance_mode = True
elif (
worker_info.state in [EdgeWorkerState.IDLE,
EdgeWorkerState.RUNNING]
and EdgeWorker.maintenance_mode
):
logger.info("Exit Maintenance mode requested!")
EdgeWorker.maintenance_mode = False
if EdgeWorker.maintenance_mode:
EdgeWorker.maintenance_comments =
worker_info.maintenance_comments
else:
EdgeWorker.maintenance_comments = None
if worker_info.state == EdgeWorkerState.SHUTDOWN_REQUEST:
logger.info("Shutdown requested!")
```
We could encapsulate the checks of the state in the WorkerInfo dataclass
through properties, then above code would be easier to read and to test also,
as then you can test the check of state directly in the WorkerInfo dataclass.
So I would add following properties in WorkerInfo dataclass:
```
def is_maintenance(self) -> bool:
return self.state == EdgeWorkerState.MAINTENANCE_REQUEST
def is_running_or_idle(self) -> bool:
return self.state in [EdgeWorkerState.IDLE, EdgeWorkerState.RUNNING]
def is_shutdown(self) -> bool:
return self.state == EdgeWorkerState.SHUTDOWN_REQUEST
```
Then you could rewrite following as below and add dedicated tests in
WorkerInfo dataclass for above properties:
```
if worker_info.is_maintenance:
logger.info("Maintenance mode requested!")
EdgeWorker.maintenance_mode = True
elif (
worker_info.is_running_or_idle
and EdgeWorker.maintenance_mode
):
logger.info("Exit Maintenance mode requested!")
EdgeWorker.maintenance_mode = False
if EdgeWorker.maintenance_mode:
EdgeWorker.maintenance_comments =
worker_info.maintenance_comments
else:
EdgeWorker.maintenance_comments = None
if worker_info.is_shutdown:
logger.info("Shutdown requested!")
```
WDYT?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
