Update tests

google · Jan 12, 2024 · 6ac474a · 6ac474a
1 parent d8622fe
commit 6ac474a
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 6 deletions.
diff --git a/tools/abel/model/entity_field.py b/tools/abel/model/entity_field.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 """Module to hold EntityField class."""
 
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Any
 
 # pylint: disable=g-importing-member
 from model.constants import BC_GUID
@@ -132,7 +132,7 @@ def FromDict(
 
   def GetSpreadsheetRowMapping(
       self, guid_to_entity_map: GuidToEntityMap
-  ) -> Dict[str, str]:
+  ) -> Dict[str, Any]:
     """Returns dictionary of spreadsheet headers to MissingField attributes."""
     missing_field_row_map = {
         VALUES: [

diff --git a/tools/abel/tests/entity_field_test.py b/tools/abel/tests/entity_field_test.py
@@ -108,7 +108,9 @@ def testMissingFieldGetSpreadsheetRowMapping(self, test_get_code):
               STRING_VALUE: TEST_MISSING_STANDARD_FIELD_NAME
             }},
             {USER_ENTERED_VALUE: {STRING_VALUE: ''}},
-            {USER_ENTERED_VALUE: {STRING_VALUE: ''}},
+            {USER_ENTERED_VALUE: {
+              STRING_VALUE: TEST_MISSING_STANDARD_FIELD_NAME
+            }},
             {USER_ENTERED_VALUE: {STRING_VALUE: TEST_REPORTING_ENTITY_CODE}},
             {USER_ENTERED_VALUE: {STRING_VALUE: TEST_REPORTING_GUID}},
             {USER_ENTERED_VALUE: {STRING_VALUE: TEST_REPORTING_ENTITY_CODE}},
@@ -133,7 +135,6 @@ def testMissingFieldGetSpreadsheetRowMapping(self, test_get_code):
     actual_row_mapping = test_missing_field.GetSpreadsheetRowMapping(
         guid_to_entity_map=test_guid_to_entity_map
     )
-
     self.assertEqual(expected_row_mapping, actual_row_mapping)
 
 

diff --git a/tools/abel/tests/export_helper_test.py b/tools/abel/tests/export_helper_test.py
@@ -195,8 +195,8 @@ def testExportBuildingConfigExportsMultiStateValueFieldStates(self):
         .get('states')
     )
 
-    self.assertIsInstance(multi_state_value_field_states.get('ON'), str)
-    self.assertEqual(multi_state_value_field_states.get('ON'), 'TRUE')
+    self.assertIsInstance(multi_state_value_field_states.get('ON'), list)
+    self.assertEqual(multi_state_value_field_states.get('ON'), ['TRUE'])
 
   def testExportBuildingConfigExportsUnitsCorrectly(self):
     exported_building_config = (