diff --git a/tests/test_newpolygons.py b/tests/test_newpolygons.py
index ffa06d0fb..0c05256c5 100644
--- a/tests/test_newpolygons.py
+++ b/tests/test_newpolygons.py
@@ -7,7 +7,7 @@
 from pathlib import Path
 from traceback import print_exception
 import warnings
-from typing import Optional, List
+from typing import Optional, List, Union
 
 from PIL import Image
 
@@ -96,7 +96,7 @@ def _test_krakencli(self, extractor_mock: Mock, *, args, force_no_legacy: bool=F
         for cl in extractor_mock.mock_calls:
             self.assertEqual(cl[2]["legacy"], expect_legacy)
 
-    def _test_ketoscli(self, *, args, expect_legacy: bool, check_exit_code: Optional[int|List[int]]=0, patching_dir="kraken.lib.dataset.recognition"):
+    def _test_ketoscli(self, *, args, expect_legacy: bool, check_exit_code: Optional[Union[int, List[int]]] = 0, patching_dir="kraken.lib.dataset.recognition"):
         """
         Base recipe for testing ketos_cli with a given polygon extraction method
         """
@@ -336,9 +336,13 @@ def test_ketoscli_pretrain_old_model(self):
             )
 
 
-    def _assertWarnsWhenTrainingArrow(
-            self, model: str, *dset: str, from_model: str|None=None, force_legacy: bool=False, 
-            expect_warning_msgs: list[str]=[], expect_not_warning_msgs: list[str]=[]):
+    def _assertWarnsWhenTrainingArrow(self,
+                                      model: str,
+                                      *dset: str,
+                                      from_model: Optional[str] = None,
+                                      force_legacy: bool = False,
+                                      expect_warning_msgs: List[str] = [],
+                                      expect_not_warning_msgs: List[str] = []):
 
         args = ['-f', 'binary', '-N', '1', '-q', 'fixed', '-o', model, *dset]
         if force_legacy:
@@ -446,4 +450,4 @@ def test_ketos_mixed_arrow_train_new(self):
                 patching_dir="kraken.lib.arrow_dataset",
             )
 
-            self._assertWarnsWhenTrainingArrow(mfp, dset, self.arrow_data, force_legacy=True, expect_warning_msgs=["WARNING Mixed legacy polygon", "WARNING Setting dataset legacy polygon status to False based on training set"], expect_not_warning_msgs=["model will be flagged to use legacy"])
\ No newline at end of file
+            self._assertWarnsWhenTrainingArrow(mfp, dset, self.arrow_data, force_legacy=True, expect_warning_msgs=["WARNING Mixed legacy polygon", "WARNING Setting dataset legacy polygon status to False based on training set"], expect_not_warning_msgs=["model will be flagged to use legacy"])