diff --git a/R/split.R b/R/split.R
index 156b3dd..f2ae63e 100644
--- a/R/split.R
+++ b/R/split.R
@@ -25,8 +25,8 @@ ensure_splits <- function(
     new_splits <- setdiff(wanted_splits, present_splits)
 
     # Extract from data_spec
-    time_col <- data_spec$pfs_col
-    event_col <- data_spec$progression_col
+    time_col <- data_spec$time_to_event_col
+    event_col <- data_spec$event_col
     pivot_time_cutoff <- data_spec$pivot_time_cutoff
     train_prop <- data_spec$train_prop
 
diff --git a/tests/testthat/test-data_spec.R b/tests/testthat/test-data_spec.R
index af6614f..7d49d6c 100644
--- a/tests/testthat/test-data_spec.R
+++ b/tests/testthat/test-data_spec.R
@@ -1,9 +1,10 @@
 test_that("prepend_to_directory() works", {
 
-  model_spec_list = list(
-    ModelSpec("mock1", fitter = zeroSum::zeroSum, directory = "mock1"),
-    ModelSpec("mock2", fitter = zeroSum::zeroSum, directory = "mock2")
-  )
+  model_spec_1 <- ModelSpec("model", "mock1", zeroSum::zeroSum, split_index = 1, time_cutoffs = 2)
+  model_spec_2 <- model_spec_1
+  model_spec_2$directory <- "mock2"
+  model_spec_list <- list(model_spec_1, model_spec_2)
+
   expect_silent(new_msl <- prepend_to_directory(model_spec_list, "prepend_me"))
   expect_equal(new_msl[[1]]$directory, "prepend_me/mock1")
   expect_equal(new_msl[[2]]$directory, "prepend_me/mock2")
diff --git a/tests/testthat/test-preprocess.R b/tests/testthat/test-preprocess.R
index 9767d64..b705a0d 100644
--- a/tests/testthat/test-preprocess.R
+++ b/tests/testthat/test-preprocess.R
@@ -8,7 +8,7 @@ test_that("discretize_ipi_features() works", {
     "ecog_performance_status" = sample(0:4, n_samples, replace = TRUE),
     "n_extranodal_sites" = sample(0:6, n_samples, replace = TRUE)
   )
-  data_spec <- DataSpec(name = "test")
+  data_spec <- DataSpec(name = "test", directory = "mock", train_prop = 0.8)
 
   pheno_tbl <- discretize_ipi_features(pheno_tbl, data_spec)
   
diff --git a/tests/testthat/test-read.R b/tests/testthat/test-read.R
index e58bd4a..6bad75f 100644
--- a/tests/testthat/test-read.R
+++ b/tests/testthat/test-read.R
@@ -3,6 +3,7 @@ test_that("read function works correctly", {
   data_spec <- DataSpec(
     name = "Schmitz et al. (2018)",
     directory = "data/schmitz",
+    train_prop = .7,
     expr_fname = "expr1.csv",
     pheno_fname = "pheno1.csv",
     patient_id_col = "patient_id",
diff --git a/tests/testthat/test-split.R b/tests/testthat/test-split.R
index d83cbf8..268f156 100644
--- a/tests/testthat/test-split.R
+++ b/tests/testthat/test-split.R
@@ -9,9 +9,10 @@ test_that("split_dataset works", {
   pheno_tbl <- generate_mock_data(
     n_samples = n_samples,
     n_genes = n_genes,
-    n_na_in_pheno = 0
+    n_na_in_pheno = 0,
+    split_index = NULL
   )[["pheno_tbl"]]
-  pheno_tbl[["split_2"]] <- "test"
+  pheno_tbl[["split_2"]] <- "test" # One split already there
   data_spec <- DataSpec(
     name = "mock",
     directory = dir,
@@ -20,18 +21,21 @@ test_that("split_dataset works", {
   )
   model_spec_1 <- ModelSpec(
     name = "dummy1",
+    directory = dir,
     fitter = zeroSum::zeroSum,
     split_index = 1:2,
     time_cutoffs = 2
   )
   model_spec_2 <- ModelSpec(
     name = "dummy2",
+    directory = dir,
     fitter = zeroSum::zeroSum,
     split_index = 1,
     time_cutoffs = 2
   )
 
   for(pivot_time_cutoff in list(NULL, 2)){
+    data_spec$pivot_time_cutoff <- pivot_time_cutoff
     new_pheno_tbl <- ensure_splits(
       pheno_tbl = pheno_tbl,
       data_spec = data_spec,