Fixed pre-existing ProgressData bugs

SEED-platform · Dec 13, 2024 · 2722c8d · 2722c8d
1 parent 30ac118
commit 2722c8d
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 16 deletions.
diff --git a/seed/data_importer/tests/test_ah_import.py b/seed/data_importer/tests/test_ah_import.py
@@ -35,8 +35,8 @@ def setUp(self):
             "data_state": DATA_STATE_MAPPING,
         }
 
-        progress_data = ProgressData(func_name="match_buildings", unique_id=self.import_file)
-        sub_progress_data = ProgressData(func_name="match_sub_progress", unique_id=self.import_file)
+        progress_data = ProgressData(func_name="match_buildings", unique_id=self.import_file.id)
+        sub_progress_data = ProgressData(func_name="match_sub_progress", unique_id=self.import_file.id)
         self.action_args = [self.import_file.id, progress_data.key, sub_progress_data.key]
 
         self.blank_result = {

diff --git a/seed/lib/progress_data/progress_data.py b/seed/lib/progress_data/progress_data.py
@@ -4,6 +4,7 @@
 """
 
 import logging
+from typing import Union
 
 from seed.decorators import get_prog_key
 from seed.utils.cache import delete_cache, get_cache, set_cache
@@ -12,11 +13,11 @@
 
 
 class ProgressData:
-    def __init__(self, func_name, unique_id, init_data=None):
+    def __init__(self, func_name, unique_id: Union[str, int], init_data=None):
         self.func_name = func_name
         self.unique_id = unique_id
         self.key = get_prog_key(func_name, unique_id)
-        self.total = None
+        self.total: Union[int, None] = None
         self.increment_by = None
 
         # Load in the initialized data, some of this may be overloaded based
@@ -30,16 +31,17 @@ def initialize(self, init_data=None):
         if init_data:
             self.data = init_data
         else:
-            self.data = {}
-            self.data["status"] = "not-started"
-            self.data["status_message"] = ""
-            self.data["progress"] = 0
-            self.data["progress_key"] = self.key
-            self.data["unique_id"] = self.unique_id
-            self.data["func_name"] = self.func_name
-            self.data["message"] = None
-            self.data["stacktrace"] = None
-            self.data["summary"] = None
+            self.data = {
+                "status": "not-started",
+                "status_message": "",
+                "progress": 0,
+                "progress_key": self.key,
+                "unique_id": self.unique_id,
+                "func_name": self.func_name,
+                "message": None,
+                "stacktrace": None,
+                "summary": None,
+            }
             self.total = None
             self.increment_by = None
 

diff --git a/seed/tests/test_derived_columns.py b/seed/tests/test_derived_columns.py
@@ -760,8 +760,8 @@ def setUp(self):
             "import_file_id": self.import_file.id,
             "data_state": DATA_STATE_MAPPING,
         }
-        progress_data = ProgressData(func_name="match_buildings", unique_id=self.import_file)
-        sub_progress_data = ProgressData(func_name="match_sub_progress", unique_id=self.import_file)
+        progress_data = ProgressData(func_name="match_buildings", unique_id=self.import_file.id)
+        sub_progress_data = ProgressData(func_name="match_sub_progress", unique_id=self.import_file.id)
         self.action_args = [self.import_file.id, progress_data.key, sub_progress_data.key]
 
         # set up factory