toolchain-utils: Fix ListMachineStates with new locking methods
ListMachineStates() in AFEManager is to print locking status of all
machines. After introducing new locking mechanism with
crrev.com/c/1793911, this function needs to be refined.
This CL is to make this function work with all three locking methods.
Examples looks like these:
Machine (Board) Status
--------------- ------
chromeos2-row9-rack7-host5.cros (snappy) locked by $USER since $T
Machine (Board) Status
--------------- ------
172.16.243.69 (??) unlocked
BUG=chromium:1006434
TEST=Called ListMachineStates with three different locking methods with
both locking and unlocking.
Change-Id: Ifa6548c0d986f5da8aa41f2864205d86856389da
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/1832642
Reviewed-by: George Burgess <gbiv@chromium.org>
Tested-by: Zhizhou Yang <zhizhouy@google.com>
Auto-Submit: Zhizhou Yang <zhizhouy@google.com>
Legacy-Commit-Queue: Commit Bot <commit-bot@chromium.org>
diff --git a/afe_lock_machine.py b/afe_lock_machine.py
index 4eae6b0..510d55f 100755
--- a/afe_lock_machine.py
+++ b/afe_lock_machine.py
@@ -10,6 +10,7 @@
from __future__ import print_function
import argparse
+import enum
import getpass
import os
import sys
@@ -46,6 +47,13 @@
"""Raised when cannot get information about lab machine from lab server."""
+class MachineType(enum.Enum):
+ """Enum class to hold machine type."""
+ AFE = 'afe'
+ LOCAL = 'local'
+ SKYLAB = 'skylab'
+
+
class AFELockManager(object):
"""Class for locking/unlocking machines vie Autotest Front End servers.
@@ -155,19 +163,45 @@
machine_list.append(r.strip())
return machine_list
- def PrintStatusHeader(self, is_lab_machine):
- """Prints the status header lines for machines.
+ def GetMachineType(self, m):
+ """Get where the machine is located.
Args:
- is_lab_machine: Boolean indicating whether to print HW Lab header or
- local machine header (different spacing).
+ m: String containing the name or ip address of machine.
+
+ Returns:
+ Value of the type in MachineType Enum.
"""
- if is_lab_machine:
- print('\nMachine (Board)\t\t\t\t\tStatus')
- print('---------------\t\t\t\t\t------\n')
+ if m in self.local_machines:
+ return MachineType.LOCAL
+ if m in self.skylab_machines:
+ return MachineType.SKYLAB
+ return MachineType.AFE
+
+ def PrintStatusHeader(self):
+ """Prints the status header lines for machines."""
+ print('\nMachine (Board)\t\t\t\t\tStatus')
+ print('---------------\t\t\t\t\t------')
+
+ def PrintStatus(self, m, state, machine_type):
+ """Prints status for a single machine.
+
+ Args:
+ m: String containing the name or ip address of machine.
+ state: A dictionary of the current state of the machine.
+ machine_type: MachineType to determine where the machine is located.
+ """
+ if state['locked']:
+ if (machine_type == MachineType.AFE and
+ m not in self.toolchain_lab_machines):
+ m += '.cros'
+ print('%s (%s)\t%slocked by %s since %s' %
+ (m, state['board'], '\t\t\t' if machine_type == MachineType.LOCAL
+ else '', state['locked_by'], state['lock_time']))
else:
- print('\nMachine (Board)\t\tStatus')
- print('---------------\t\t------\n')
+ print(
+ '%s (%s)\t\t%sunlocked' % (m, state['board'], '\t\t' if
+ machine_type == MachineType.LOCAL else ''))
def AddMachineToLocal(self, machine):
"""Adds a machine to local machine list.
@@ -198,34 +232,11 @@
the current AFELockManager's list of machines. Normally obtained by
calling AFELockManager::GetMachineStates.
"""
- local_machines = []
- printed_hdr = False
+ self.PrintStatusHeader()
for m in machine_states:
- cros_name = m + '.cros'
- if (m in self.toolchain_lab_machines or
- cros_name in self.toolchain_lab_machines):
- name = m if m in self.toolchain_lab_machines else cros_name
- if not printed_hdr:
- self.PrintStatusHeader(True)
- printed_hdr = True
- state = machine_states[m]
- if state['locked']:
- print('%s (%s)\tlocked by %s since %s' %
- (name, state['board'], state['locked_by'], state['lock_time']))
- else:
- print('%s (%s)\tunlocked' % (name, state['board']))
- else:
- local_machines.append(m)
-
- if local_machines:
- self.PrintStatusHeader(False)
- for m in local_machines:
- state = machine_states[m]
- if state['locked']:
- print('%s (%s)\tlocked by %s since %s' %
- (m, state['board'], state['locked_by'], state['lock_time']))
- else:
- print('%s (%s)\tunlocked' % (m, state['board']))
+ machine_type = self.GetMachineType(m)
+ state = machine_states[m]
+ self.PrintStatus(m, state, machine_type)
def UpdateLockInAFE(self, should_lock_machine, machine):
"""Calls an AFE server to lock/unlock a machine.
@@ -317,23 +328,21 @@
for m in self.machines:
# TODO(zhizhouy): Handling exceptions with more details when locking
# doesn't succeed.
- machine_type = 'afe'
- if m in self.skylab_machines:
+ machine_type = self.GetMachineType(m)
+ if machine_type == MachineType.SKYLAB:
ret = self.UpdateLockInSkylab(lock_machines, m)
- machine_type = 'skylab'
- elif m in self.local_machines:
+ elif machine_type == MachineType.LOCAL:
ret = self.UpdateFileLock(lock_machines, m)
- machine_type = 'local'
else:
ret = self.UpdateLockInAFE(lock_machines, m)
if ret:
self.logger.LogOutput(
- '%s %s machine succeeded: %s.' % (action, machine_type, m))
+ '%s %s machine succeeded: %s.' % (action, machine_type.value, m))
updated_machines.append(m)
else:
self.logger.LogOutput(
- '%s %s machine failed: %s.' % (action, machine_type, m))
+ '%s %s machine failed: %s.' % (action, machine_type.value, m))
self.machines = updated_machines
return updated_machines
@@ -377,8 +386,10 @@
'(%s).' % k)
self._InternalRemoveMachine(k)
- if state['locked'] and 'locked_by' in state and \
- state['locked_by'] != self.user:
+ # TODO(zhizhouy): Skylab doesn't support host info such as locked_by.
+ # Need to update this when skylab supports it.
+ if (state['locked'] and state['locked_by'] and
+ state['locked_by'] != self.user):
raise DontOwnLock('Attempt to unlock machine (%s) locked by someone '
'else (%s).' % (k, state['locked_by']))
elif cmd == 'lock':
@@ -417,7 +428,13 @@
# as afe does. We need to get more info such as locked_by when skylab
# supports that.
if m in self.local_machines or m in self.skylab_machines:
- machine_list[m] = {'locked': 0 if cmd == 'lock' else 1}
+ values = {
+ 'locked': 0 if cmd == 'lock' else 1,
+ 'board': '??',
+ 'locked_by': '',
+ 'lock_time': ''
+ }
+ machine_list[m] = values
else:
# For autotest machines, we use afe APIs to get locking info.
mod_host = m.split('.')[0]