From 3fda40ac54d78bb231cee7c4e2943b588561591a Mon Sep 17 00:00:00 2001
From: Ang Li <angli@google.com>
Date: Thu, 22 Aug 2024 10:30:44 -0700
Subject: [PATCH] Include k-v pairs from device config as `user_added_info` by
 default. (#935)

Since these k-v pairs are added by users through the device config, by definition these are `user_added_info`. So they should be part of the `user_added_info` field of `device_info`.

Because they are user added, we need to guard against arbitrary values that may break yaml reporting downstream.
---
 mobly/controllers/android_device.py           | 14 ++++++++++-
 .../mobly/controllers/android_device_test.py  | 24 +++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/mobly/controllers/android_device.py b/mobly/controllers/android_device.py
index 866f50b5..b147da42 100644
--- a/mobly/controllers/android_device.py
+++ b/mobly/controllers/android_device.py
@@ -139,8 +139,19 @@ def get_info(ads):
 
   Returns:
     A list of dict, each representing info for an AndroidDevice objects.
+    Everything in this dict should be yaml serializable.
   """
-  return [ad.device_info for ad in ads]
+  infos = []
+  # The values of user_added_info can be arbitrary types, so we shall sanitize
+  # them here to ensure they are yaml serializable.
+  for ad in ads:
+    device_info = ad.device_info
+    user_added_info = {
+        k: str(v) for (k, v) in device_info['user_added_info'].items()
+    }
+    device_info['user_added_info'] = user_added_info
+    infos.append(device_info)
+  return infos
 
 
 def _validate_device_existence(serials):
@@ -908,6 +919,7 @@ def load_config(self, config):
             % (k, getattr(self, k)),
         )
       setattr(self, k, v)
+      self.add_device_info(k, v)
 
   def root_adb(self):
     """Change adb to root mode for this device if allowed.
diff --git a/tests/mobly/controllers/android_device_test.py b/tests/mobly/controllers/android_device_test.py
index 68fcfe4e..557af5f7 100755
--- a/tests/mobly/controllers/android_device_test.py
+++ b/tests/mobly/controllers/android_device_test.py
@@ -153,6 +153,29 @@ def test_create_with_no_valid_config(self):
     with self.assertRaisesRegex(android_device.Error, expected_msg):
       android_device.create([1])
 
+  @mock.patch(
+      'mobly.controllers.android_device_lib.adb.AdbProxy',
+      return_value=mock_android_device.MockAdbProxy(1),
+  )
+  @mock.patch(
+      'mobly.controllers.android_device_lib.fastboot.FastbootProxy',
+      return_value=mock_android_device.MockFastbootProxy(1),
+  )
+  @mock.patch('mobly.utils.create_dir')
+  def test_get_info(self, create_dir_mock, FastbootProxy, MockAdbProxy):
+    mock_serial = '1'
+    ad = android_device.AndroidDevice(serial=mock_serial)
+    example_user_object = mock_android_device.MockAdbProxy('magic')
+    # Add an arbitrary object as a device info
+    ad.add_device_info('user_stuff', example_user_object)
+    info = android_device.get_info([ad])[0]
+    self.assertEqual(info['serial'], mock_serial)
+    self.assertTrue(info['build_info'])
+    # User added values should be normalized to strings.
+    self.assertEqual(
+        info['user_added_info']['user_stuff'], str(example_user_object)
+    )
+
   @mock.patch('mobly.controllers.android_device.list_fastboot_devices')
   @mock.patch('mobly.controllers.android_device.list_adb_devices')
   @mock.patch('mobly.controllers.android_device.list_adb_devices_by_usb_id')
@@ -393,6 +416,7 @@ def test_AndroidDevice_load_config(
     self.assertEqual(ad.space, 'the final frontier')
     self.assertEqual(ad.number, 1)
     self.assertEqual(ad.debug_tag, 'my_tag')
+    self.assertEqual(ad.device_info['user_added_info']['debug_tag'], 'my_tag')
 
   @mock.patch(
       'mobly.controllers.android_device_lib.adb.AdbProxy',