BigQuery: use _get_sub_prop helper so missing load stats don't raise. (…

…#6269)
googleapis · Oct 19, 2018 · 333c52a · 333c52a
1 parent a74dc33
commit 333c52a
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 17 deletions.
diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py
@@ -1374,9 +1374,9 @@ def input_file_bytes(self):
         :returns: the count (None until set from the server).
         :raises: ValueError for invalid value types.
         """
-        statistics = self._properties.get('statistics')
-        if statistics is not None:
-            return int(statistics['load']['inputFileBytes'])
+        return _helpers._int_or_none(_helpers._get_sub_prop(
+                self._properties, ['statistics', 'load', 'inputFileBytes']
+            ))
 
     @property
     def input_files(self):
@@ -1385,9 +1385,9 @@ def input_files(self):
         :rtype: int, or ``NoneType``
         :returns: the count (None until set from the server).
         """
-        statistics = self._properties.get('statistics')
-        if statistics is not None:
-            return int(statistics['load']['inputFiles'])
+        return _helpers._int_or_none(_helpers._get_sub_prop(
+                self._properties, ['statistics', 'load', 'inputFiles']
+            ))
 
     @property
     def output_bytes(self):
@@ -1396,9 +1396,9 @@ def output_bytes(self):
         :rtype: int, or ``NoneType``
         :returns: the count (None until set from the server).
         """
-        statistics = self._properties.get('statistics')
-        if statistics is not None:
-            return int(statistics['load']['outputBytes'])
+        return _helpers._int_or_none(_helpers._get_sub_prop(
+                self._properties, ['statistics', 'load', 'outputBytes']
+            ))
 
     @property
     def output_rows(self):
@@ -1407,9 +1407,9 @@ def output_rows(self):
         :rtype: int, or ``NoneType``
         :returns: the count (None until set from the server).
         """
-        statistics = self._properties.get('statistics')
-        if statistics is not None:
-            return int(statistics['load']['outputRows'])
+        return _helpers._int_or_none(_helpers._get_sub_prop(
+                self._properties, ['statistics', 'load', 'outputRows']
+            ))
 
     def to_api_repr(self):
         """Generate a resource for :meth:`_begin`."""

diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py
@@ -1985,11 +1985,6 @@ def test_props_set_by_server(self):
         statistics['creationTime'] = _millis(CREATED)
         statistics['startTime'] = _millis(STARTED)
         statistics['endTime'] = _millis(ENDED)
-        load_stats = statistics['load'] = {}
-        load_stats['inputFileBytes'] = 12345
-        load_stats['inputFiles'] = 1
-        load_stats['outputBytes'] = 23456
-        load_stats['outputRows'] = 345
 
         self.assertEqual(job.etag, 'ETAG')
         self.assertEqual(job.self_link, URL)
@@ -1999,6 +1994,15 @@ def test_props_set_by_server(self):
         self.assertEqual(job.started, STARTED)
         self.assertEqual(job.ended, ENDED)
 
+        # running jobs have no load stats not yet set.
+        self.assertIsNone(job.output_bytes)
+
+        load_stats = statistics['load'] = {}
+        load_stats['inputFileBytes'] = 12345
+        load_stats['inputFiles'] = 1
+        load_stats['outputBytes'] = 23456
+        load_stats['outputRows'] = 345
+
         self.assertEqual(job.input_file_bytes, 12345)
         self.assertEqual(job.input_files, 1)
         self.assertEqual(job.output_bytes, 23456)