diff --git a/gfe_integration_tests/sdram_edge_tests/test_constant_internal/sdram_splitter_internal.py b/gfe_integration_tests/sdram_edge_tests/test_constant_internal/sdram_splitter_internal.py
index 94efa3c9..4fac689a 100644
--- a/gfe_integration_tests/sdram_edge_tests/test_constant_internal/sdram_splitter_internal.py
+++ b/gfe_integration_tests/sdram_edge_tests/test_constant_internal/sdram_splitter_internal.py
@@ -47,11 +47,11 @@ def _post_vertex(self):
         return self.__post_vertex
 
     @overrides(AbstractSplitterCommon.get_out_going_vertices)
-    def get_out_going_vertices(self, partition_id) -> List[SDRAMMachineVertex]:
+    def get_out_going_vertices(self, partition_id: str) -> List[SDRAMMachineVertex]:
         return [self._pre_vertex]
 
     @overrides(AbstractSplitterCommon.get_in_coming_vertices)
-    def get_in_coming_vertices(self, partition_id) -> list[SDRAMMachineVertex]:
+    def get_in_coming_vertices(self, partition_id: str) -> list[SDRAMMachineVertex]:
         return [self._post_vertex]
 
     def create_machine_vertices(self, chip_counter):
diff --git a/gfe_integration_tests/sdram_edge_tests/test_multi_dest_internal/sdram_splitter.py b/gfe_integration_tests/sdram_edge_tests/test_multi_dest_internal/sdram_splitter.py
index 5ad9d3b0..61180fb3 100644
--- a/gfe_integration_tests/sdram_edge_tests/test_multi_dest_internal/sdram_splitter.py
+++ b/gfe_integration_tests/sdram_edge_tests/test_multi_dest_internal/sdram_splitter.py
@@ -47,11 +47,11 @@ def __post_vertex(self):
         return self.__post_vertex
 
     @overrides(AbstractSplitterCommon.get_out_going_vertices)
-    def get_out_going_vertices(self, partition_id) -> List[SDRAMMachineVertex]:
+    def get_out_going_vertices(self, partition_id: str) -> List[SDRAMMachineVertex]:
         return self._post_vertices
 
     @overrides(AbstractSplitterCommon.get_in_coming_vertices)
-    def get_in_coming_vertices(self, partition_id) -> SDRAMMachineVertex:
+    def get_in_coming_vertices(self, partition_id: str) -> SDRAMMachineVertex:
         return [self._pre_vertex]
 
     def create_machine_vertices(self, chip_counter):
diff --git a/gfe_integration_tests/sdram_edge_tests/test_multi_src_internal/sdram_splitter.py b/gfe_integration_tests/sdram_edge_tests/test_multi_src_internal/sdram_splitter.py
index fb0313b5..ebf92f0e 100644
--- a/gfe_integration_tests/sdram_edge_tests/test_multi_src_internal/sdram_splitter.py
+++ b/gfe_integration_tests/sdram_edge_tests/test_multi_src_internal/sdram_splitter.py
@@ -43,12 +43,12 @@ def _post_vertex(self):
 
     @overrides(AbstractSplitterCommon.get_out_going_vertices)
     def get_out_going_vertices(
-            self, partition_id) -> SourceSegmentedSDRAMMachinePartition:
+            self, partition_id: str) -> SourceSegmentedSDRAMMachinePartition:
         return [self._post_vertex]
 
     @overrides(AbstractSplitterCommon.get_in_coming_vertices)
     def get_in_coming_vertices(
-            self, partition_id) -> List[SourceSegmentedSDRAMMachinePartition]:
+            self, partition_id: str) -> List[SourceSegmentedSDRAMMachinePartition]:
         return self._pre_vertices
 
     def create_machine_vertices(self, chip_counter):