diff --git a/test/reference/ddr3_init.h b/test/reference/ddr3_init.h
index 2e6ce47..c5c9139 100644
--- a/test/reference/ddr3_init.h
+++ b/test/reference/ddr3_init.h
@@ -34,6 +34,8 @@
 #define SDRAM_PHY_MODULES 8
 #define SDRAM_PHY_DELAYS 32
 #define SDRAM_PHY_BITSLIPS 8
+#define SDRAM_PHY_DDR3
+#define SDRAM_PHY_SUPPORTED_MEMORY 0x0000000040000000ULL
 
 void cdelay(int i);
 
diff --git a/test/reference/ddr4_init.h b/test/reference/ddr4_init.h
index 7163058..b6ddbf0 100644
--- a/test/reference/ddr4_init.h
+++ b/test/reference/ddr4_init.h
@@ -33,6 +33,8 @@
 #define SDRAM_PHY_MODULES 8
 #define SDRAM_PHY_DELAYS 512
 #define SDRAM_PHY_BITSLIPS 8
+#define SDRAM_PHY_DDR4
+#define SDRAM_PHY_SUPPORTED_MEMORY 0x0000000080000000ULL
 
 void cdelay(int i);
 
diff --git a/test/reference/sdr_init.h b/test/reference/sdr_init.h
index 51aff84..4e0947b 100644
--- a/test/reference/sdr_init.h
+++ b/test/reference/sdr_init.h
@@ -27,6 +27,8 @@
 #define SDRAM_PHY_WRPHASE 0
 #define SDRAM_PHY_DQ_DQS_RATIO 8
 #define SDRAM_PHY_MODULES 2
+#define SDRAM_PHY_SDR
+#define SDRAM_PHY_SUPPORTED_MEMORY 0x0000000002000000ULL
 
 void cdelay(int i);
 
diff --git a/test/test_init.py b/test/test_init.py
index 84ccf0c..f9ac127 100644
--- a/test/test_init.py
+++ b/test/test_init.py
@@ -29,8 +29,15 @@ class TestInit(unittest.TestCase):
     def test_sdr(self):
         from litex_boards.targets.scarabhardware_minispartan6 import BaseSoC
         soc       = BaseSoC()
-        c_header  = get_sdram_phy_c_header(soc.sdram.controller.settings.phy, soc.sdram.controller.settings.timing)
-        py_header = get_sdram_phy_py_header(soc.sdram.controller.settings.phy, soc.sdram.controller.settings.timing)
+        c_header  = get_sdram_phy_c_header(
+            phy_settings    = soc.sdram.controller.settings.phy,
+            timing_settings = soc.sdram.controller.settings.timing,
+            geom_settings   = soc.sdram.controller.settings.geom,
+        )
+        py_header = get_sdram_phy_py_header(
+            phy_settings    = soc.sdram.controller.settings.phy,
+            timing_settings = soc.sdram.controller.settings.timing
+        )
         #update_c_reference(c_header, "sdr_init.h")
         compare_with_reference(self, c_header, "sdr_init.h")
         compare_with_reference(self, py_header, "sdr_init.py")
@@ -38,8 +45,15 @@ class TestInit(unittest.TestCase):
     def test_ddr3(self):
         from litex_boards.targets.xilinx_kc705 import BaseSoC
         soc       = BaseSoC()
-        c_header  = get_sdram_phy_c_header(soc.sdram.controller.settings.phy, soc.sdram.controller.settings.timing)
-        py_header = get_sdram_phy_py_header(soc.sdram.controller.settings.phy, soc.sdram.controller.settings.timing)
+        c_header  = get_sdram_phy_c_header(
+            phy_settings    = soc.sdram.controller.settings.phy,
+            timing_settings = soc.sdram.controller.settings.timing,
+            geom_settings   = soc.sdram.controller.settings.geom,
+        )
+        py_header = get_sdram_phy_py_header(
+            phy_settings    = soc.sdram.controller.settings.phy,
+            timing_settings = soc.sdram.controller.settings.timing
+        )
         #update_c_reference(c_header, "ddr3_init.h")
         compare_with_reference(self, c_header, "ddr3_init.h")
         compare_with_reference(self, py_header, "ddr3_init.py")
@@ -47,8 +61,15 @@ class TestInit(unittest.TestCase):
     def test_ddr4(self):
         from litex_boards.targets.xilinx_kcu105 import BaseSoC
         soc       = BaseSoC(max_sdram_size=0x4000000)
-        c_header  = get_sdram_phy_c_header(soc.sdram.controller.settings.phy, soc.sdram.controller.settings.timing)
-        py_header = get_sdram_phy_py_header(soc.sdram.controller.settings.phy, soc.sdram.controller.settings.timing)
+        c_header  = get_sdram_phy_c_header(
+            phy_settings    = soc.sdram.controller.settings.phy,
+            timing_settings = soc.sdram.controller.settings.timing,
+            geom_settings   = soc.sdram.controller.settings.geom,
+        )
+        py_header = get_sdram_phy_py_header(
+            phy_settings    = soc.sdram.controller.settings.phy,
+            timing_settings = soc.sdram.controller.settings.timing
+        )
         #update_c_reference(c_header, "ddr4_init.h")
         compare_with_reference(self, c_header, "ddr4_init.h")
         compare_with_reference(self, py_header, "ddr4_init.py")