diff --git a/CITATION.cff b/CITATION.cff
index 9c457910..6352874f 100644
--- a/CITATION.cff
+++ b/CITATION.cff
@@ -9,7 +9,7 @@ type: software
 license: MIT
 title: 'bpmodels: Simulating and Analysing Transmission Chain Statistics using Branching
   Process Models'
-version: 0.3.0
+version: 0.3.1
 abstract: Provides methods to simulate and analyse the size and length of branching
   processes with an arbitrary offspring distribution. These can be used, for example,
   to analyse the distribution of chain sizes or length of infectious disease outbreaks,
@@ -249,3 +249,15 @@ references:
     given-names: Saralees
     email: saralees.nadarajah@manchester.ac.uk
   year: '2023'
+- type: software
+  title: checkmate
+  abstract: 'checkmate: Fast and Versatile Argument Checks'
+  notes: Imports
+  url: https://mllg.github.io/checkmate/
+  repository: https://CRAN.R-project.org/package=checkmate
+  authors:
+  - family-names: Lang
+    given-names: Michel
+    email: michellang@gmail.com
+    orcid: https://orcid.org/0000-0001-9754-0393
+  year: '2023'
diff --git a/DESCRIPTION b/DESCRIPTION
index ec22a9a8..36ceeab7 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,7 +1,7 @@
 Package: bpmodels
 Title: Simulating and Analysing Transmission Chain Statistics using Branching Process
     Models
-Version: 0.3.0
+Version: 0.3.1
 Authors@R: c(
     person("James M.", "Azam", , "james.azam@lshtm.ac.uk", role = c("aut", "cre"),
            comment = c(ORCID = "https://orcid.org/0000-0001-5782-7330")),
@@ -48,3 +48,5 @@ LazyData: true
 Roxygen: list(markdown = TRUE)
 RoxygenNote: 7.2.3
 Language: en-GB
+Imports: 
+    checkmate
diff --git a/NEWS.md b/NEWS.md
index c73baebe..a77841fe 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,3 +1,10 @@
+# bpmodels 0.3.1
+
+## Unit tests and input validation
+
+* The following internal functions now have input validation: `rborel()`, `dborel()`, `complementary_logprob()`, and `rnbinom_mean_disp()`.
+* Code coverage has been improved with more tests on the following functions: `rborel()`, `dborel()`, `chain_sim()`, `rnbinom_mean_disp()`, `complementary_logprob()`, `rgen_length()`, and `rbinom_size()`.
+
 # bpmodels 0.3.0
 
 ## Website
diff --git a/R/borel.r b/R/borel.r
index 7ce580ab..9ec85985 100644
--- a/R/borel.r
+++ b/R/borel.r
@@ -1,12 +1,17 @@
 ##' Density of the Borel distribution
 ##'
-##' @param x vector of integers.
-##' @param mu mu parameter.
+##' @param x vector of quantiles; integer.
+##' @param mu mu parameter (the poisson mean); non-negative.
 ##' @param log logical; if TRUE, probabilities p are given as log(p).
 ##' @return probability mass.
 ##' @author Sebastian Funk
 dborel <- function(x, mu, log = FALSE) {
-  if (x < 1) stop("'x' must be greater than 0")
+  checkmate::assert_numeric(
+    x, lower = 1, upper = Inf
+  )
+  checkmate::assert_number(
+    mu, lower = 0, finite = TRUE, na.ok = FALSE
+  )
   ld <- -mu * x + (x - 1) * log(mu * x) - lgamma(x + 1)
   if (!log) ld <- exp(ld)
   return(ld)
@@ -16,11 +21,17 @@ dborel <- function(x, mu, log = FALSE) {
 ##'
 ##' Random numbers are generated by simulating from a Poisson branching process
 ##' @param n number of random variates to generate.
-##' @param mu mu parameter.
+##' @param mu mu parameter (the Poisson mean).
 ##' @param infinite any number to treat as infinite; simulations will be stopped
 ##'     if this number is reached
 ##' @return vector of random numbers
 ##' @author Sebastian Funk
 rborel <- function(n, mu, infinite = Inf) {
+  checkmate::assert_number(
+    n, lower = 1, finite = TRUE, na.ok = FALSE
+  )
+  checkmate::assert_number(
+    mu, lower = 0, finite = TRUE, na.ok = FALSE
+  )
   chain_sim(n, "pois", "size", infinite = infinite, lambda = mu)
 }
diff --git a/R/utils.r b/R/utils.r
index c4d135a5..76397c50 100644
--- a/R/utils.r
+++ b/R/utils.r
@@ -6,6 +6,9 @@
 #' @author Sebastian Funk
 #' @keywords internal
 complementary_logprob <- function(x) {
+  checkmate::assert_numeric(
+    x, lower = -Inf, upper = 0
+  )
   tryCatch(log1p(-sum(exp(x))), error = function(e) -Inf)
 }
 
@@ -43,14 +46,23 @@ rgen_length <- function(n, x, prob) {
 #' Negative binomial random numbers parametrized
 #' in terms of mean and dispersion coefficient
 #' @param n number of samples to draw
-#' @param mn mean of distribution
-#' @param disp dispersion coefficient (var/mean)
+#' @param mn mean of distribution; Must be > 0.
+#' @param disp dispersion coefficient (var/mean); Must be > 1.
 #' @return vector containing the random numbers
 #' @author Flavio Finger
 #' @export
 #' @examples
 #' rnbinom_mean_disp(n = 5, mn = 4, disp = 2)
 rnbinom_mean_disp <- function(n, mn, disp) {
+  checkmate::assert_number(
+    n, lower = 1, finite = TRUE, na.ok = FALSE
+  )
+  checkmate::assert_number(
+    disp, lower = 1, finite = TRUE, na.ok = FALSE
+  )
+  checkmate::assert_number(
+    mn, lower = 1E-100, finite = TRUE, na.ok = FALSE
+  )
   size <- mn / (disp - 1)
   stats::rnbinom(n, size = size, mu = mn)
 }
diff --git a/man/dborel.Rd b/man/dborel.Rd
index 14d269d0..7871bf46 100644
--- a/man/dborel.Rd
+++ b/man/dborel.Rd
@@ -7,9 +7,9 @@
 dborel(x, mu, log = FALSE)
 }
 \arguments{
-\item{x}{vector of integers.}
+\item{x}{vector of quantiles; integer.}
 
-\item{mu}{mu parameter.}
+\item{mu}{mu parameter (the poisson mean); non-negative.}
 
 \item{log}{logical; if TRUE, probabilities p are given as log(p).}
 }
diff --git a/man/rborel.Rd b/man/rborel.Rd
index 5562f7e1..0303b4e0 100644
--- a/man/rborel.Rd
+++ b/man/rborel.Rd
@@ -9,7 +9,7 @@ rborel(n, mu, infinite = Inf)
 \arguments{
 \item{n}{number of random variates to generate.}
 
-\item{mu}{mu parameter.}
+\item{mu}{mu parameter (the Poisson mean).}
 
 \item{infinite}{any number to treat as infinite; simulations will be stopped
 if this number is reached}
diff --git a/man/rnbinom_mean_disp.Rd b/man/rnbinom_mean_disp.Rd
index 698836d6..fb9df0aa 100644
--- a/man/rnbinom_mean_disp.Rd
+++ b/man/rnbinom_mean_disp.Rd
@@ -10,9 +10,9 @@ rnbinom_mean_disp(n, mn, disp)
 \arguments{
 \item{n}{number of samples to draw}
 
-\item{mn}{mean of distribution}
+\item{mn}{mean of distribution; Must be > 0.}
 
-\item{disp}{dispersion coefficient (var/mean)}
+\item{disp}{dispersion coefficient (var/mean); Must be > 1.}
 }
 \value{
 vector containing the random numbers
diff --git a/tests/testthat/test-borel.r b/tests/testthat/test-borel.r
new file mode 100644
index 00000000..2c2829b5
--- /dev/null
+++ b/tests/testthat/test-borel.r
@@ -0,0 +1,28 @@
+test_that("We can calculate probabilities and sample", {
+  expect_gt(dborel(1, 0.5), 0)
+  expect_identical(dborel(1, 0.5, log = TRUE), -0.5)
+  expect_length(rborel(2, 0.9), 2)
+})
+
+test_that("Errors are thrown", {
+  expect_error(
+    dborel(x = 0, mu = 0.5),
+    "is not >= 1"
+  )
+  expect_error(
+    dborel(x = 1, mu = -0.5),
+    "is not >= 0"
+  )
+  expect_error(
+    dborel(x = 1, mu = Inf),
+    "Must be finite"
+  )
+  expect_error(
+    rborel(n = 0, mu = -0.5),
+    "is not >= 1"
+  )
+  expect_error(
+    rborel(n = 0, mu = Inf),
+    "is not >= 1"
+  )
+})
diff --git a/tests/testthat/test-utils.R b/tests/testthat/test-utils.R
new file mode 100644
index 00000000..d96e78bf
--- /dev/null
+++ b/tests/testthat/test-utils.R
@@ -0,0 +1,60 @@
+test_that("Util functions work", {
+  expect_length(rnbinom_mean_disp(n = 5, mn = 4, disp = 2), 5)
+  expect_length(rgen_length(n = 1, x = c(1, 2, 3), prob = 0.3), 3)
+  expect_length(rbinom_size(n = 1, x = c(1, 2, 3), prob = 0.3), 3)
+  expect_identical(complementary_logprob(x = 0), -Inf)
+  expect_identical(complementary_logprob(x = -Inf), 0)
+  expect_lt(complementary_logprob(x = -0.1), 0)
+})
+
+test_that("Errors are thrown", {
+  # Checks on 'disp' argument
+  expect_error(
+    rnbinom_mean_disp(n = 5, mn = 4, disp = 0.9),
+    "is not >= 1"
+  )
+  expect_error(
+    rnbinom_mean_disp(n = 5, mn = 4, disp = NA),
+    "May not be NA"
+  )
+  expect_error(
+    rnbinom_mean_disp(n = 5, mn = 4, disp = Inf),
+    "Must be finite"
+  )
+  # Checks on 'n' argument
+  expect_error(
+    rnbinom_mean_disp(n = 0, mn = 4, disp = 2),
+    "is not >= 1"
+  )
+  expect_error(
+    rnbinom_mean_disp(n = NA, mn = 4, disp = 2),
+    "May not be NA"
+  )
+  expect_error(
+    rnbinom_mean_disp(n = Inf, mn = 4, disp = 2),
+    "Must be finite"
+  )
+  # Checks on 'mn' argument
+  expect_error(
+    rnbinom_mean_disp(n = 2, mn = 0, disp = 2)
+  )
+  expect_error(
+    rnbinom_mean_disp(n = 2, mn = NA, disp = 2),
+    "May not be NA"
+  )
+  expect_error(
+    rnbinom_mean_disp(n = 2, mn = Inf, disp = 2),
+    "Must be finite"
+  )
+})
+
+test_that("Errors are thrown", {
+  expect_error(
+    complementary_logprob(0.1),
+    "is not <= 0"
+  )
+  expect_error(
+    complementary_logprob(Inf),
+    "is not <= 0"
+  )
+})
diff --git a/tests/testthat/tests-borel.r b/tests/testthat/tests-borel.r
deleted file mode 100644
index e17512e1..00000000
--- a/tests/testthat/tests-borel.r
+++ /dev/null
@@ -1,9 +0,0 @@
-test_that("We can calculate probabilities and sample", {
-  expect_gt(dborel(1, 0.5), 0)
-  expect_identical(dborel(1, 0.5, log = TRUE), -0.5)
-  expect_length(rborel(2, 0.9), 2)
-})
-
-test_that("Errors are thrown", {
-  expect_error(dborel(0, 0.5), "greater than 0")
-})
diff --git a/tests/testthat/tests-sim.r b/tests/testthat/tests-sim.r
index bb053a43..c284b323 100644
--- a/tests/testthat/tests-sim.r
+++ b/tests/testthat/tests-sim.r
@@ -1,49 +1,133 @@
 test_that("Chains can be simulated", {
-  expect_length(chain_sim(n = 2, "pois", lambda = 0.5), 2)
-  expect_length(chain_sim(n = 10, "pois", "length", lambda = 0.9), 10)
+  set.seed(12)
+  tf <- 3
+  chain_sim_test_df <- chain_sim(
+    n = 2,
+    offspring = "pois",
+    stat = "size",
+    lambda = 0.9,
+    tree = TRUE,
+    serial = function(n) {
+      rlnorm(n, meanlog = 0.58, sdlog = 1.58)
+    },
+    tf = tf
+  )
+  # Check that all the simulated times are less than tf
+  expect_true(
+    all(
+      chain_sim_test_df$time < tf
+    )
+  )
+  # Other checks
+  expect_length(
+    chain_sim(
+      n = 2,
+      offspring = "pois",
+      lambda = 0.5
+    ),
+    2
+  )
+  expect_length(
+    chain_sim(
+      n = 10,
+      offspring = "pois",
+      stat = "length",
+      lambda = 0.9
+    ),
+    10
+  )
   expect_s3_class(
-    chain_sim(n = 10,
-              "pois",
-              lambda = 2,
-              tree = TRUE,
-              infinite = 10
-              ),
+    chain_sim(
+      n = 10,
+      offspring = "pois",
+      lambda = 2,
+      tree = TRUE,
+      infinite = 10
+    ),
     "data.frame"
+  )
+  expect_false(
+    any(
+      is.finite(
+        chain_sim(
+          n = 2,
+          offspring = "pois",
+          stat = "length",
+          lambda = 0.5,
+          infinite = 1
+        )
+      )
     )
-  expect_false(any(is.finite(chain_sim(
-    n = 2, "pois", "length", lambda = 0.5,
-    infinite = 1
-  ))))
-  expect_no_error(chain_sim(
-    n = 2, offspring = "pois", "size", lambda = 0.9,
-    tree = TRUE
-  ))
+  )
+  expect_no_error(
+    chain_sim(
+      n = 2,
+      offspring = "pois",
+      stat = "size",
+      lambda = 0.9,
+      tree = TRUE
+    )
+  )
 })
 
 test_that("Errors are thrown", {
-  expect_error(chain_sim(n = 2, "dummy"), "does not exist")
-  expect_error(chain_sim(n = 2, "lnorm", meanlog = log(1.6)), "integer")
   expect_error(
-    chain_sim(n = 2, offspring = pois, "length", lambda = 0.9),
+    chain_sim(
+      n = 2,
+      "dummy"
+    ),
+    "does not exist"
+  )
+  expect_error(
+    chain_sim(
+      n = 2,
+      offspring = "lnorm",
+      meanlog = log(1.6)
+    ),
+    "integer"
+  )
+  expect_error(
+    chain_sim(
+      n = 2,
+      offspring = pois,
+      stat = "length",
+      lambda = 0.9
+    ),
     "not found"
   )
-  expect_error(chain_sim(
-    n = 2, offspring = "pois", "size", lambda = 0.9,
-    serial = c(1, 2), "must be a function"
-  ))
   expect_error(
-    chain_sim(n = 2, offspring = c(1, 2), "length", lambda = 0.9),
+    chain_sim(
+      n = 2,
+      offspring = "pois",
+      stat = "size",
+      lambda = 0.9,
+      serial = c(1, 2)
+    ),
+    "must be a function"
+  )
+  expect_error(
+    chain_sim(
+      n = 2,
+      offspring = c(1, 2),
+      stat = "length",
+      lambda = 0.9
+    ),
     "not a character string"
   )
   expect_error(
-    chain_sim(n = 2, offspring = list(1, 2), "length", lambda = 0.9),
+    chain_sim(
+      n = 2,
+      offspring = list(1, 2),
+      stat = "length",
+      lambda = 0.9
+    ),
     "not a character string"
   )
   expect_error(
     chain_sim(
       n = 2,
       offspring = "pois",
-      "size",
+      stat = "size",
       lambda = 0.9,
       tf = 5,
       tree = FALSE
@@ -54,30 +138,30 @@ test_that("Errors are thrown", {
 
 test_that("Chains can be simulated", {
   expect_s3_class(
-      chain_sim_susc(
-        "pois",
-        mn_offspring = 2,
-        serial = function(x) 3,
-        pop = 100
-      ),
-      "data.frame"
+    chain_sim_susc(
+      offspring = "pois",
+      mn_offspring = 2,
+      serial = function(x) 3,
+      pop = 100
+    ),
+    "data.frame"
   )
 
   expect_s3_class(
-      chain_sim_susc(
-        "nbinom",
-        mn_offspring = 2,
-        disp_offspring = 1.5,
-        serial = function(x) 3,
-        pop = 100
-      ),
-      "data.frame"
+    chain_sim_susc(
+      offspring = "nbinom",
+      mn_offspring = 2,
+      disp_offspring = 1.5,
+      serial = function(x) 3,
+      pop = 100
+    ),
+    "data.frame"
   )
 
   expect_identical(
     nrow(
       chain_sim_susc(
-        "pois",
+        offspring = "pois",
         mn_offspring = 2,
         serial = function(x) 3,
         pop = 1
@@ -89,7 +173,7 @@ test_that("Chains can be simulated", {
   expect_identical(
     nrow(
       chain_sim_susc(
-        "pois",
+        offspring = "pois",
         mn_offspring = 100,
         tf = 2,
         serial = function(x) 3,
@@ -102,7 +186,7 @@ test_that("Chains can be simulated", {
   expect_identical(
     nrow(
       chain_sim_susc(
-        "pois",
+        offspring = "pois",
         mn_offspring = 100,
         serial = function(x) 3,
         pop = 999,
@@ -116,7 +200,7 @@ test_that("Chains can be simulated", {
 test_that("Errors are thrown", {
   expect_error(
     chain_sim_susc(
-      "dummy",
+      offspring = "dummy",
       mn_offspring = 3,
       serial = function(x) 3,
       pop = 100
@@ -125,20 +209,21 @@ test_that("Errors are thrown", {
   )
   expect_error(
     chain_sim_susc(
-      "nbinom",
+      offspring = "nbinom",
       mn_offspring = 3,
       disp_offspring = 1,
       serial = function(x) 3,
       pop = 100
     ),
-    paste("Offspring distribution 'nbinom'",
-          "requires argument 'disp_offspring' > 1.",
-          "Use 'pois' if there is no overdispersion."
-          )
+    paste(
+      "Offspring distribution 'nbinom'",
+      "requires argument 'disp_offspring' > 1.",
+      "Use 'pois' if there is no overdispersion."
     )
+  )
   expect_error(
     chain_sim_susc(
-      "nbinom",
+      offspring = "nbinom",
       mn_offspring = 3,
       serial = function(x) 3,
       pop = 100
@@ -150,7 +235,7 @@ test_that("Errors are thrown", {
 test_that("warnings work as expected", {
   expect_warning(
     chain_sim_susc(
-      "pois",
+      offspring = "pois",
       mn_offspring = 3,
       disp_offspring = 1,
       serial = function(x) 3,
@@ -162,14 +247,15 @@ test_that("warnings work as expected", {
     chain_sim(
       n = 2,
       offspring = "pois",
-      "size",
+      stat = "size",
       lambda = 0.9,
       serial = function(x) rpois(x, 0.9),
       tree = FALSE
     ),
-    sprintf("%s %s",
-            "`serial` can't be used with `tree = FALSE`;",
-            "Setting `tree = TRUE` internally."
+    sprintf(
+      "%s %s",
+      "`serial` can't be used with `tree = FALSE`;",
+      "Setting `tree = TRUE` internally."
     )
   )
 })