fix: fixed serialization of total_items

open-contracting · May 15, 2021 · 055ff65 · 055ff65
1 parent 21b7131
commit 055ff65
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/spoonbill/stats.py b/spoonbill/stats.py
@@ -32,6 +32,7 @@ class DataPreprocessor:
     :param combined_tables: List of tables with data from different locations
     :param tables: Do not parse schema and use this tables data
     :param table_threshold: Maximum array length before system recommends it to separated to child table
+    :param total_items: Total objects processed
     """
 
     def __init__(
@@ -41,6 +42,7 @@ def __init__(
         combined_tables: Mapping[str, List] = None,
         tables: Mapping[str, Table] = None,
         table_threshold=TABLE_THRESHOLD,
+        total_items=0,
         header_separator="/",
     ):
         self.schema = schema
@@ -50,7 +52,7 @@ def __init__(
         self.table_threshold = table_threshold
 
         self.header_separator = header_separator
-        self.total_items = 0
+        self.total_items = total_items
         self.current_table = None
 
         self._lookup_cache = {}
@@ -297,6 +299,7 @@ def dump(self):
             "header_separator": self.header_separator,
             "tables": {name: table.dump() for name, table in self.tables.items()},
             "table_threshold": self.table_threshold,
+            "total_items": self.total_items,
         }
 
     @classmethod
@@ -312,6 +315,7 @@ def restore(cls, data):
                 "combined_tables": data["combined_tables"],
                 "header_separator": data["header_separator"],
                 "table_threshold": data["table_threshold"],
+                "total_items": data["total_items"],
             }
         except KeyError as e:
             LOGGER.error(_("Failed to restore from malformed data. Missing {} attribute").format(e))

diff --git a/tests/test_stats.py b/tests/test_stats.py
@@ -145,7 +145,16 @@ def test_dump_restore(spec, releases, tmpdir):
     spec2 = DataPreprocessor.restore(data)
     for name, table in spec.tables.items():
         assert table == spec2.tables[name]
-
+    for key in (
+        "schema",
+        "root_tables",
+        "combined_tables",
+        "header_separator",
+        "tables",
+        "table_threshold",
+        "total_items",
+    ):
+        assert key in spec2.__dict__
     with pytest.raises(ValueError, match="Unable to restore"):
         del data["schema"]
         spec2 = DataPreprocessor.restore(data)