diff --git a/EpiAware/README.md b/EpiAware/README.md
index d5fee9898..e2a2edf2e 100644
--- a/EpiAware/README.md
+++ b/EpiAware/README.md
@@ -9,63 +9,68 @@
 - Solid lines indicate implemented features/analysis.
 - Dashed lines indicate planned features/analysis.
 
-## Proposed `EpiAware` model diagram
+## Current `EpiAware` model diagram
 ```mermaid
 flowchart LR
 
-    A["Underlying dists.
-and specify length of sims
----------------------
-EpiData"]
+A["Underlying GI
+Bijector"]
+
+EpiModel["AbstractEpiModel
+----------------------
+Choice of target
+for latent process:
 
-    B["Choice of target
-for latent process
----------------------
 DirectInfections
     ExpGrowthRate
     Renewal"]
 
-C["Observational Data
+InitModel["Priors for
+initial scale of incidence"]
+
+DataW[Data wrangling and QC]
+
+
+ObsData["Observational Data
 ---------------------
 Obs. cases y_t"]
-D["Latent processes
+
+LatentProcPriors["Latent process priors"]
+
+LatentProc["AbstractLatentProcess
+---------------------
+RandomWalkLatentProcess"]
+
+ObsModelPriors["Observation model priors
+choice of delayed obs. model"]
+
+ObsModel["AbstractObservationModel
 ---------------------
-random_walk"]
+DelayObservations"]
+
 E["Turing model constructor
 ---------------------
 make_epi_inference_model"]
-F["Latent Process priors
----------------------
-default_rw_priors"]
+
 G[Posterior draws]
 H[Posterior checking]
 I[Post-processing]
-DataW[Data wrangling and QC]
-J["Observation models
----------------------
-delay_observations"]
-K["Observation model priors
----------------------
-default_delay_obs_priors"]
-ObservationModel["ObservationModel
----------------------
-delay_observations_model"]
-LatentProcess["LatentProcess
----------------------
-random_walk_process"]
 
-A --> EpiModel
-B --> EpiModel
+
+
+A --> EpiData
+EpiData --> EpiModel
+InitModel --> EpiModel
 EpiModel -->E
-C-->E
-D-->LatentProcess
-F-->LatentProcess
-J-->ObservationModel
-K-->ObservationModel
-LatentProcess-->E
-ObservationModel-->E
+ObsData-->E
+DataW-.->ObsData
+LatentProcPriors-->LatentProc
+LatentProc-->E
+ObsModelPriors-->ObsModel
+ObsModel-->E
+
+
 E-->|sample...NUTS...| G
 G-.->H
 H-.->I
-DataW-.->C
 ```
diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl
index b38ff340d..35a5e7eec 100644
--- a/EpiAware/src/EpiAware.jl
+++ b/EpiAware/src/EpiAware.jl
@@ -32,21 +32,18 @@ using Distributions,
       DataFramesMeta
 
 # Exported utilities
-export create_discrete_pmf, default_rw_priors, default_delay_obs_priors,
-       default_initialisation_prior, spread_draws
+export create_discrete_pmf, spread_draws
 
 # Exported types
 export EpiData, Renewal, ExpGrowthRate, DirectInfections
 
 # Exported Turing model constructors
-export make_epi_inference_model, delay_observations_model, random_walk_process,
-       initialize_incidence
+export make_epi_inference_model
 
 include("epimodel.jl")
 include("utilities.jl")
 include("latent-processes.jl")
 include("observation-processes.jl")
-include("initialisation.jl")
 include("models.jl")
 
 end
diff --git a/EpiAware/src/epimodel.jl b/EpiAware/src/epimodel.jl
index 67929d821..adb993575 100644
--- a/EpiAware/src/epimodel.jl
+++ b/EpiAware/src/epimodel.jl
@@ -2,91 +2,117 @@ abstract type AbstractEpiModel end
 
 struct EpiData{T <: Real, F <: Function}
     gen_int::Vector{T}
-    delay_int::Vector{T}
-    delay_kernel::SparseMatrixCSC{T, Integer}
-    cluster_coeff::T
     len_gen_int::Integer
-    len_delay_int::Integer
-    time_horizon::Integer
     transformation::F
 
     #Inner constructors for EpiData object
-    function EpiData(
-            gen_int,
-            delay_int,
-            cluster_coeff,
-            time_horizon::Integer,
-            transformation::Function
-    )
+    function EpiData(gen_int,
+            transformation::Function)
         @assert all(gen_int .>= 0) "Generation interval must be non-negative"
-        @assert all(delay_int .>= 0) "Delay interval must be non-negative"
         @assert sum(gen_int)≈1 "Generation interval must sum to 1"
-        @assert sum(delay_int)≈1 "Delay interval must sum to 1"
 
-        K = generate_observation_kernel(delay_int, time_horizon)
-
-        new{eltype(gen_int), typeof(transformation)}(
-            gen_int,
-            delay_int,
-            K,
-            cluster_coeff,
+        new{eltype(gen_int), typeof(transformation)}(gen_int,
             length(gen_int),
-            length(delay_int),
-            time_horizon,
-            transformation
-        )
+            transformation)
     end
 
-    function EpiData(
-            gen_distribution::ContinuousDistribution,
-            delay_distribution::ContinuousDistribution,
-            cluster_coeff,
-            time_horizon::Integer;
+    function EpiData(gen_distribution::ContinuousDistribution;
             D_gen,
-            D_delay,
             Δd = 1.0,
-            transformation::Function = exp
-    )
+            transformation::Function = exp)
         gen_int = create_discrete_pmf(gen_distribution, Δd = Δd, D = D_gen) |>
                   p -> p[2:end] ./ sum(p[2:end])
-        delay_int = create_discrete_pmf(delay_distribution, Δd = Δd, D = D_delay)
 
-        return EpiData(gen_int, delay_int, cluster_coeff, time_horizon, transformation)
+        return EpiData(gen_int, transformation)
     end
 end
 
-struct DirectInfections <: AbstractEpiModel
+struct DirectInfections{S <: Sampleable} <: AbstractEpiModel
     data::EpiData
+    initialisation_prior::S
 end
 
-function (epimodel::DirectInfections)(_It, init)
-    epimodel.data.transformation.(init .+ _It)
+struct ExpGrowthRate{S <: Sampleable} <: AbstractEpiModel
+    data::EpiData
+    initialisation_prior::S
 end
 
-struct ExpGrowthRate <: AbstractEpiModel
+struct Renewal{S <: Sampleable} <: AbstractEpiModel
     data::EpiData
+    initialisation_prior::S
 end
 
-function (epimodel::ExpGrowthRate)(rt, init)
-    init .+ cumsum(rt) .|> exp
+"""
+    function (epimodel::Renewal)(recent_incidence, Rt)
+
+Compute new incidence based on recent incidence and Rt.
+
+This is a callable function on `Renewal` structs, that encodes new incidence prediction
+given recent incidence and Rt according to basic renewal process.
+
+```math
+I_t = R_t \\sum_{i=1}^{n-1} I_{t-i} g_i
+```
+
+where `I_t` is the new incidence, `R_t` is the reproduction number, `I_{t-i}` is the recent incidence
+and `g_i` is the generation interval.
+
+
+# Arguments
+- `recent_incidence`: Array of recent incidence values.
+- `Rt`: Reproduction number.
+
+# Returns
+- Tuple containing the updated incidence array and the new incidence value.
+
+"""
+function (epimodel::Renewal)(recent_incidence, Rt)
+    new_incidence = Rt * dot(recent_incidence, epimodel.data.gen_int)
+    return ([new_incidence; recent_incidence[1:(epimodel.data.len_gen_int - 1)]],
+        new_incidence)
 end
 
-struct Renewal <: AbstractEpiModel
-    data::EpiData
+function generate_latent_infs(epimodel::AbstractEpiModel, latent_process)
+    @info "No concrete implementation for `generate_latent_infs` is defined."
+    return nothing
 end
 
-function (epimodel::Renewal)(_Rt, init)
-    I₀ = epimodel.data.transformation(init)
+@model function generate_latent_infs(epimodel::DirectInfections, _It)
+    init_incidence ~ epimodel.initialisation_prior
+    return epimodel.data.transformation.(init_incidence .+ _It)
+end
+
+@model function generate_latent_infs(epimodel::ExpGrowthRate, rt)
+    init_incidence ~ epimodel.initialisation_prior
+    return exp.(init_incidence .+ cumsum(rt))
+end
+
+"""
+    generate_latent_infs(epimodel::Renewal, _Rt)
+
+`Turing` model constructor for latent infections using the `Renewal` object `epimodel` and time-varying unconstrained reproduction number `_Rt`.
+
+`generate_latent_infs` creates a `Turing` model for sampling latent infections with given unconstrainted
+reproduction number `_Rt` but random initial incidence scale. The initial incidence pre-time one is given as
+a scale on top of an exponential growing process with exponential growth rate given by `R_to_r`applied to the
+first value of `Rt`.
+
+# Arguments
+- `epimodel::Renewal`: The epidemiological model.
+- `_Rt`: Time-varying unconstrained (e.g. log-) reproduction number.
+
+# Returns
+- `I_t`: Array of latent infections over time.
+
+"""
+@model function generate_latent_infs(epimodel::Renewal, _Rt)
+    init_incidence ~ epimodel.initialisation_prior
+    I₀ = epimodel.data.transformation(init_incidence)
     Rt = epimodel.data.transformation.(_Rt)
 
     r_approx = R_to_r(Rt[1], epimodel)
     init = I₀ * [exp(-r_approx * t) for t in 0:(epimodel.data.len_gen_int - 1)]
 
-    function generate_infs(recent_incidence, Rt)
-        new_incidence = Rt * dot(recent_incidence, epimodel.data.gen_int)
-        [new_incidence; recent_incidence[1:(epimodel.data.len_gen_int - 1)]], new_incidence
-    end
-
-    I_t, _ = scan(generate_infs, init, Rt)
+    I_t, _ = scan(epimodel, init, Rt)
     return I_t
 end
diff --git a/EpiAware/src/initialisation.jl b/EpiAware/src/initialisation.jl
deleted file mode 100644
index 6bf975933..000000000
--- a/EpiAware/src/initialisation.jl
+++ /dev/null
@@ -1,30 +0,0 @@
-"""
-    default_initialisation_prior()
-
-Constructs a default initialisation prior for the model.
-
-# Returns
-`NamedTuple` with the following fields:
-- `I0_prior`: A standard normal distribution representing the prior for the initial infected population.
-
-"""
-function default_initialisation_prior()
-    (; I0_prior = Normal(),)
-end
-
-"""
-    initialize_incidence(; I0_prior)
-
-Initialize the incidence of the disease in unconstrained domain.
-
-# Arguments
-- `I0_prior::Distribution`: Prior distribution for the initial incidence.
-
-# Returns
-- `_I0`: Unconstrained initial incidence value.
-
-"""
-@model function initialize_incidence(; I0_prior::Distribution)
-    _I0 ~ I0_prior
-    return _I0
-end
diff --git a/EpiAware/src/latent-processes.jl b/EpiAware/src/latent-processes.jl
index eabf12ec8..89f78d0ae 100644
--- a/EpiAware/src/latent-processes.jl
+++ b/EpiAware/src/latent-processes.jl
@@ -1,51 +1,30 @@
+abstract type AbstractLatentProcess end
+
+struct RandomWalkLatentProcess{D <: Sampleable, S <: Sampleable} <: AbstractLatentProcess
+    init_prior::D
+    var_prior::S
+end
+
 function default_rw_priors()
-    return (
-        :var_RW_prior => truncated(Normal(0.0, 0.05), 0.0, Inf),
-        :init_rw_value_prior => Normal()
-    ) |> Dict
+    return (:var_RW_prior => truncated(Normal(0.0, 0.05), 0.0, Inf),
+        :init_rw_value_prior => Normal()) |> Dict
+end
+
+function generate_latent_process(latent_process::AbstractLatentProcess, n)
+    @info "No concrete implementation for generate_latent_process is defined."
+    return nothing
 end
 
-@model function random_walk(n; var_RW_prior, init_rw_value_prior)
+@model function generate_latent_process(latent_process::RandomWalkLatentProcess, n)
     ϵ_t ~ MvNormal(ones(n))
-    σ²_RW ~ var_RW_prior
-    init ~ init_rw_value_prior
+    σ²_RW ~ latent_process.var_prior
+    rw_init ~ latent_process.init_prior
     σ_RW = sqrt(σ²_RW)
     rw = Vector{eltype(ϵ_t)}(undef, n)
 
-    rw[1] = σ_RW * ϵ_t[1]
+    rw[1] = rw_init + σ_RW * ϵ_t[1]
     for t in 2:n
         rw[t] = rw[t - 1] + σ_RW * ϵ_t[t]
     end
-    return rw, init, (; σ_RW,)
-end
-
-"""
-    struct LatentProcess{F<:Function}
-
-A struct representing a latent process with its priors.
-
-# Fields
-- `latent_process`: The latent process function for a `Turing` model.
-- `latent_process_priors`: NamedTuple containing the priors for the latent process.
-
-"""
-struct LatentProcess{F <: Function, D <: Distribution}
-    latent_process::F
-    latent_process_priors::Dict{Symbol, D}
-end
-
-"""
-    random_walk_process(; latent_process_priors = default_rw_priors())
-
-Create a `LatentProcess` struct reflecting a random walk process with optional priors.
-
-# Arguments
-- `latent_process_priors`: Optional priors for the random walk process.
-
-# Returns
-- `LatentProcess`: A random walk process.
-
-"""
-function random_walk_process(; latent_process_priors = default_rw_priors())
-    LatentProcess(random_walk, latent_process_priors)
+    return rw, (; σ_RW, rw_init)
 end
diff --git a/EpiAware/src/models.jl b/EpiAware/src/models.jl
index 8b3e73064..73a54e23a 100644
--- a/EpiAware/src/models.jl
+++ b/EpiAware/src/models.jl
@@ -1,34 +1,27 @@
-@model function make_epi_inference_model(
-        y_t,
+@model function make_epi_inference_model(y_t,
+        time_steps;
         epimodel::AbstractEpiModel,
-        latent_process_obj::LatentProcess,
-        observation_process_obj::ObservationModel;
-        pos_shift = 1e-6
-)
+        latent_process_model::AbstractLatentProcess,
+        observation_model::AbstractObservationModel,
+        pos_shift = 1e-6)
     #Latent process
-    time_steps = epimodel.data.time_horizon
-    @submodel latent_process, init, latent_process_aux = latent_process_obj.latent_process(
-        time_steps;
-        latent_process_obj.latent_process_priors...
-    )
+    @submodel latent_process, latent_process_aux = generate_latent_process(
+        latent_process_model,
+        time_steps)
 
     #Transform into infections
-    I_t = epimodel(latent_process, init)
+    @submodel I_t = generate_latent_infs(epimodel, latent_process)
 
     #Predictive distribution of ascerted cases
-    @submodel generated_y_t, generated_y_t_aux = observation_process_obj.observation_model(
+    @submodel generated_y_t, generated_y_t_aux = generate_observations(observation_model,
         y_t,
-        I_t,
-        epimodel::AbstractEpiModel;
-        pos_shift = pos_shift,
-        observation_process_obj.observation_model_priors...
-    )
+        I_t;
+        pos_shift = pos_shift)
 
     #Generate quantities
     return (;
         generated_y_t,
         I_t,
         latent_process,
-        process_aux = merge(latent_process_aux, generated_y_t_aux)
-    )
+        process_aux = merge(latent_process_aux, generated_y_t_aux))
 end
diff --git a/EpiAware/src/observation-processes.jl b/EpiAware/src/observation-processes.jl
index 023c67dcc..fe2bab10f 100644
--- a/EpiAware/src/observation-processes.jl
+++ b/EpiAware/src/observation-processes.jl
@@ -1,19 +1,53 @@
+abstract type AbstractObservationModel end
+
+struct DelayObservations{T <: AbstractFloat, S <: Sampleable} <: AbstractObservationModel
+    delay_kernel::SparseMatrixCSC{T, Integer}
+    neg_bin_cluster_factor_prior::S
+
+    function DelayObservations(delay_int,
+            time_horizon,
+            neg_bin_cluster_factor_prior)
+        @assert all(delay_int .>= 0) "Delay interval must be non-negative"
+        @assert sum(delay_int)≈1 "Delay interval must sum to 1"
+
+        K = generate_observation_kernel(delay_int, time_horizon)
+
+        new{eltype(K), typeof(neg_bin_cluster_factor_prior)}(K,
+            neg_bin_cluster_factor_prior)
+    end
+
+    function DelayObservations(;
+            delay_distribution::ContinuousDistribution,
+            time_horizon::Integer,
+            neg_bin_cluster_factor_prior::Sampleable,
+            D_delay,
+            Δd = 1.0)
+        delay_int = create_discrete_pmf(delay_distribution; Δd = Δd, D = D_delay)
+        return DelayObservations(delay_int, time_horizon, neg_bin_cluster_factor_prior)
+    end
+end
+
 function default_delay_obs_priors()
     return (:neg_bin_cluster_factor_prior => Gamma(3, 0.05 / 3),) |> Dict
 end
 
-@model function delay_observations(
+function generate_observations(observation_model::AbstractObservationModel,
+        y_t,
+        I_t;
+        pos_shift)
+    @info "No concrete implementation for generate_observations is defined."
+    return nothing
+end
+
+@model function generate_observations(observation_model::DelayObservations,
         y_t,
-        I_t,
-        epimodel::AbstractEpiModel;
-        neg_bin_cluster_factor_prior,
-        pos_shift
-)
+        I_t;
+        pos_shift)
     #Parameters
-    neg_bin_cluster_factor ~ neg_bin_cluster_factor_prior
+    neg_bin_cluster_factor ~ observation_model.neg_bin_cluster_factor_prior
 
     #Predictive distribution
-    case_pred_dists = (epimodel.data.delay_kernel * I_t) .+ pos_shift .|>
+    case_pred_dists = (observation_model.delay_kernel * I_t) .+ pos_shift .|>
                       μ -> mean_cc_neg_bin(μ, neg_bin_cluster_factor)
 
     #Likelihood
@@ -21,34 +55,3 @@ end
 
     return y_t, (; neg_bin_cluster_factor,)
 end
-
-"""
-    struct ObservationModel{F <: Function, D<:Distribution}
-
-A struct representing an observation model.
-
-# Fields
-- `observation_model`: The observation model function.
-- `observation_model_priors`: A dictionary of prior distributions for the observation model parameters.
-
-"""
-struct ObservationModel{F <: Function, D <: Distribution}
-    observation_model::F
-    observation_model_priors::Dict{Symbol, D}
-end
-
-"""
-    delay_observations_model(; latent_process_priors = default_rw_priors())
-
-Create an `ObservationModel` struct reflecting a delayed observation process with optional priors.
-
-# Arguments
-- `latent_process_priors`: Optional priors for the delayed observation process.
-
-# Returns
-- `ObservationModel`: An observation model with delayed observations.
-
-"""
-function delay_observations_model(; observation_model_priors = default_delay_obs_priors())
-    ObservationModel(delay_observations, observation_model_priors)
-end
diff --git a/EpiAware/src/utilities.jl b/EpiAware/src/utilities.jl
index 5a2f9d035..35c2c4c1f 100644
--- a/EpiAware/src/utilities.jl
+++ b/EpiAware/src/utilities.jl
@@ -17,7 +17,7 @@ value. This is similar to the JAX function `jax.lax.scan`.
 - `ys`: An array containing the result values of applying `f` to each element of `xs`.
 - `carry`: The final accumulator value.
 """
-function scan(f::Function, init, xs::Vector{T}) where {T <: Union{Integer, AbstractFloat}}
+function scan(f, init, xs::Vector{T}) where {T <: Union{Integer, AbstractFloat}}
     carry = init
     ys = similar(xs)
     for (i, x) in enumerate(xs)
@@ -52,13 +52,11 @@ Raises:
 - `AssertionError` if `Δd` is not positive.
 - `AssertionError` if `D` is not greater than `Δd`.
 """
-function create_discrete_pmf(
-        dist::Distribution,
+function create_discrete_pmf(dist::Distribution,
         ::Val{:single_censored};
         primary_approximation_point = 0.5,
         Δd = 1.0,
-        D
-)
+        D)
     @assert minimum(dist)>=0.0 "Distribution must be non-negative"
     @assert Δd>0.0 "Δd must be positive"
     @assert D>Δd "D must be greater than Δd"
@@ -201,7 +199,7 @@ Compute the mean-cluster factor negative binomial distribution.
 A `NegativeBinomial` distribution object.
 """
 function mean_cc_neg_bin(μ, α)
-    ex_σ² = α * μ^2
+    ex_σ² = (α * μ^2) + 1e-6
     p = μ / (μ + ex_σ² + 1e-6)
     r = μ^2 / ex_σ²
     return NegativeBinomial(r, p)
diff --git a/EpiAware/test/predictive_checking/discretized_pmfs.jl b/EpiAware/test/predictive_checking/discretized_pmfs.jl
index da6e71ef6..f12b881d1 100644
--- a/EpiAware/test/predictive_checking/discretized_pmfs.jl
+++ b/EpiAware/test/predictive_checking/discretized_pmfs.jl
@@ -72,16 +72,14 @@ plt1 = let
     pmf1 = create_discrete_pmf(cont_dist, Val(:single_censored); Δd = Δd, D = D)
     pmf2 = create_discrete_pmf(cont_dist; Δd = Δd, D = D)
 
-    bar(
-        ts,
+    bar(ts,
         [pmf1;; pmf2],
         fillalpha = 0.5,
         lw = 0,
         title = "Discrete PMF with Δd = 1 day",
         label = ["Single censoring (midpoint primary)" "Double Censoring"],
         xlabel = "Days",
-        ylabel = "Probability"
-    )
+        ylabel = "Probability")
 end
 savefig(plt1, joinpath(@__DIR__(), "assets/", "discrete_pmf_daily.png"))
 
@@ -93,15 +91,13 @@ plt2 = let
     pmf1 = create_discrete_pmf(cont_dist, Val(:single_censored); Δd = Δd, D = D)
     pmf2 = create_discrete_pmf(cont_dist; Δd = Δd, D = D)
 
-    bar(
-        ts,
+    bar(ts,
         [pmf1;; pmf2],
         fillalpha = 0.5,
         lw = 0,
         title = "Discrete PMF with Δd = 1 hour",
         label = ["Single censoring (midpoint primary)" "Double Censoring"],
         xlabel = "Days",
-        ylabel = "Probability"
-    )
+        ylabel = "Probability")
 end
 savefig(plt2, joinpath(@__DIR__(), "assets/", "discrete_pmf_hourly.png"))
diff --git a/EpiAware/test/predictive_checking/fast_approx_for_r.jl b/EpiAware/test/predictive_checking/fast_approx_for_r.jl
index a36f89dd0..a6cd82f5f 100644
--- a/EpiAware/test/predictive_checking/fast_approx_for_r.jl
+++ b/EpiAware/test/predictive_checking/fast_approx_for_r.jl
@@ -65,8 +65,7 @@ errors = mapreduce(hcat, doubling_times) do T_2
     end
 end
 
-plot(
-    idxs,
+plot(idxs,
     errors,
     yscale = :log10,
     xlabel = "Newton steps",
@@ -74,5 +73,4 @@ plot(
     title = "Fast approximation for r",
     lab = ["T_2 = 1.0" "T_2 = 3.5" "T_2 = 7.0" "T_2 = 14.0"],
     yticks = [0.0, 1e-15, 1e-10, 1e-5, 1e0] |> x -> (x .+ jitter, string.(x)),
-    xticks = 0:2:10
-)
+    xticks = 0:2:10)
diff --git a/EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl b/EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl
index 3de96a657..1170a634a 100644
--- a/EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl
+++ b/EpiAware/test/predictive_checking/toy_model_log_infs_RW.jl
@@ -77,22 +77,14 @@ Random.seed!(0)
 
 - Medium length generation interval distribution.
 - Median 2 day, std 4.3 day delay distribution.
-- 100 days of simulations
 =#
 
 truth_GI = Gamma(2, 5)
-truth_delay = LogNormal(2.0, 1.0)
-neg_bin_cluster_factor = 0.05
-time_horizon = 100
+model_data = EpiData(truth_GI,
+    D_gen = 10.0)
 
-model_data = EpiData(
-    truth_GI,
-    truth_delay,
-    neg_bin_cluster_factor,
-    time_horizon,
-    D_gen = 10.0,
-    D_delay = 10.0
-)
+log_I0_prior = Normal(0.0, 1.0)
+epimodel = DirectInfections(model_data, log_I0_prior)
 
 #=
 ## Define the data generating process
@@ -100,17 +92,24 @@ model_data = EpiData(
 In this case we use the `DirectInfections` model.
 =#
 
-toy_log_infs = DirectInfections(model_data)
-rwp = random_walk_process()
+rwp = EpiAware.RandomWalkLatentProcess(Normal(),
+    truncated(Normal(0.0, 0.01), 0.0, 0.5))
 obs_mdl = delay_observations_model()
 
+#Define the observation model - no delay model
+time_horizon = 100
+obs_model = EpiAware.DelayObservations([1.0],
+    time_horizon,
+    truncated(Gamma(5, 0.05 / 5), 1e-3, 1.0))
+
 #=
 ## Generate a `Turing` `Model`
 We don't have observed data, so we use `missing` value for `y_t`.
 =#
 
-log_infs_model = make_epi_inference_model(
-    missing, toy_log_infs, rwp, obs_mdl; pos_shift = 1e-6)
+log_infs_model = make_epi_inference_model(missing, time_horizon, ; epimodel = epimodel,
+    latent_process_model = rwp, observation_model = obs_model,
+    pos_shift = 1e-6)
 
 #=
 ## Sample from the model
@@ -126,13 +125,11 @@ cond_toy = fix(log_infs_model, (init = log(1.0), σ²_RW = 0.1))
 random_epidemic = rand(cond_toy)
 gen = generated_quantities(cond_toy, random_epidemic)
 
-plot(
-    gen.I_t,
+plot(gen.I_t,
     label = "I_t",
     xlabel = "Time",
     ylabel = "Infections",
-    title = "Generated Infections"
-)
+    title = "Generated Infections")
 scatter!(random_epidemic.y_t, lab = "generated cases")
 
 #=
@@ -143,16 +140,15 @@ We treat the generated data as observed data and attempt to infer underlying inf
 
 truth_data = random_epidemic.y_t
 
-model = make_epi_inference_model(truth_data, toy_log_infs, rwp, obs_mdl; pos_shift = 1e-6)
-
-@time chn = sample(
-    model,
+model = make_epi_inference_model(truth_data, time_horizon, ; epimodel = epimodel,
+    latent_process_model = rwp, observation_model = obs_model,
+    pos_shift = 1e-6)
+@time chn = sample(model,
     NUTS(; adtype = AutoReverseDiff(true)),
     MCMCThreads(),
     250,
     4;
-    drop_warmup = true
-)
+    drop_warmup = true)
 
 #=
 ## Postior predictive checking
@@ -165,14 +161,12 @@ predicted_y_t = mapreduce(hcat, generated_quantities(log_infs_model, chn)) do ge
 end
 
 plot(predicted_y_t, c = :grey, alpha = 0.05, lab = "")
-scatter!(
-    truth_data,
+scatter!(truth_data,
     lab = "Observed cases",
     xlabel = "Time",
     ylabel = "Cases",
     title = "Posterior Predictive Checking",
-    ylims = (-0.5, maximum(truth_data) * 2.5)
-)
+    ylims = (-0.5, maximum(truth_data) * 2.5))
 
 #=
 ## Underlying inferred infections
@@ -183,14 +177,12 @@ predicted_I_t = mapreduce(hcat, generated_quantities(log_infs_model, chn)) do ge
 end
 
 plot(predicted_I_t, c = :grey, alpha = 0.05, lab = "")
-scatter!(
-    gen.I_t,
+scatter!(gen.I_t,
     lab = "Actual infections",
     xlabel = "Time",
     ylabel = "Unobserved Infections",
     title = "Posterior Predictive Checking",
-    ylims = (-0.5, maximum(gen.I_t) * 1.5)
-)
+    ylims = (-0.5, maximum(gen.I_t) * 1.5))
 
 #=
 ## Outputing the MCMC chain
diff --git a/EpiAware/test/prior_predictive_checking/ppc-latent-processes.jl b/EpiAware/test/prior_predictive_checking/ppc-latent-processes.jl
index 462b267d8..67fc8f333 100644
--- a/EpiAware/test/prior_predictive_checking/ppc-latent-processes.jl
+++ b/EpiAware/test/prior_predictive_checking/ppc-latent-processes.jl
@@ -23,43 +23,35 @@ end
 theoretical_std = [t * latent_process_priors.var_RW_prior.untruncated.σ * sqrt(2) / sqrt(π)
                    for t in 1:n] .|> sqrt
 
-plt_ppc_rw = plot(
-    sampled_walks, lab = "", ylabel = "RW", xlabel = "t", c = :grey, alpha = 0.1)
-plot!(
-    plt_ppc_rw,
+plt_ppc_rw = plot(sampled_walks, lab = "", ylabel = "RW", xlabel = "t", c = :grey,
+    alpha = 0.1)
+plot!(plt_ppc_rw,
     zeros(n),
     lw = 2,
     c = :red,
     lab = "Theoretical 3 sigma spread",
     ribbon = 3 * theoretical_std,
-    fillalpha = 0.2
-)
+    fillalpha = 0.2)
 
-σ_hist = histogram(
-    prior_chn[:σ²_RW],
+σ_hist = histogram(prior_chn[:σ²_RW],
     norm = :pdf,
     lab = "",
     ylabel = "Density",
     xlabel = "σ²_RW",
     c = :grey,
-    alpha = 0.5
-)
-plot!(
-    σ_hist,
+    alpha = 0.5)
+plot!(σ_hist,
     latent_process_priors.var_RW_prior,
     lw = 2,
     c = :red,
     alpha = 0.5,
     lab = "Prior",
-    bins = 100
-)
+    bins = 100)
 
-plt_rw = plot(
-    plt_ppc_rw,
+plt_rw = plot(plt_ppc_rw,
     σ_hist,
     layout = (1, 2),
     size = (800, 400),
     left_margin = 3mm,
-    bottom_margin = 3mm
-)
+    bottom_margin = 3mm)
 savefig(plt_rw, joinpath(@__DIR__(), "assets", "ppc_rw.png"))
diff --git a/EpiAware/test/test_epimodel.jl b/EpiAware/test/test_epimodel.jl
index c975ef350..b8d5f4a7e 100644
--- a/EpiAware/test/test_epimodel.jl
+++ b/EpiAware/test/test_epimodel.jl
@@ -1,23 +1,13 @@
 
 @testitem "EpiData constructor" begin
     gen_int = [0.2, 0.3, 0.5]
-    delay_int = [0.1, 0.4, 0.5]
-    cluster_coeff = 0.8
-    time_horizon = 10
     transformation = exp
 
-    data = EpiData(gen_int, delay_int, cluster_coeff, time_horizon, transformation)
+    data = EpiData(gen_int, transformation)
 
     @test length(data.gen_int) == 3
-    @test length(data.delay_int) == 3
-    @test data.cluster_coeff == 0.8
     @test data.len_gen_int == 3
-    @test data.len_delay_int == 3
-
     @test sum(data.gen_int) ≈ 1
-    @test sum(data.delay_int) ≈ 1
-
-    @test size(data.delay_kernel) == (time_horizon, time_horizon)
     @test data.transformation(0.0) == 1.0
 end
 
@@ -25,42 +15,29 @@ end
     using Distributions
 
     gen_distribution = Uniform(0.0, 10.0)
-    delay_distribution = Exponential(1.0)
     cluster_coeff = 0.8
     time_horizon = 10
     D_gen = 10.0
-    D_delay = 10.0
     Δd = 1.0
 
-    data = EpiData(
-        gen_distribution,
-        delay_distribution,
-        cluster_coeff,
-        time_horizon;
-        D_gen = 10.0,
-        D_delay = 10.0
-    )
+    data = EpiData(gen_distribution;
+        D_gen = 10.0)
 
-    @test data.cluster_coeff == 0.8
     @test data.len_gen_int == Int64(D_gen / Δd) - 1
-    @test data.len_delay_int == Int64(D_delay / Δd)
 
     @test sum(data.gen_int) ≈ 1
-    @test sum(data.delay_int) ≈ 1
-
-    @test size(data.delay_kernel) == (time_horizon, time_horizon)
 end
 
 @testitem "Renewal function: internal generate infs" begin
-    using LinearAlgebra
+    using LinearAlgebra, Distributions
     gen_int = [0.2, 0.3, 0.5]
     delay_int = [0.1, 0.4, 0.5]
     cluster_coeff = 0.8
     time_horizon = 10
     transformation = exp
 
-    data = EpiData(gen_int, delay_int, cluster_coeff, time_horizon, transformation)
-    epimodel = Renewal(data)
+    data = EpiData(gen_int, transformation)
+    epimodel = Renewal(data, Normal())
 
     function generate_infs(recent_incidence, Rt)
         new_incidence = Rt * dot(recent_incidence, epimodel.data.gen_int)
@@ -77,36 +54,111 @@ end
     @test generate_infs(recent_incidence, Rt) == expected_output
 end
 
-@testitem "ExpGrowthRate function" begin
+@testitem "generate_latent_infs dispatched on ExpGrowthRate" begin
+    using Distributions, Turing, HypothesisTests, DynamicPPL
     gen_int = [0.2, 0.3, 0.5]
-    delay_int = [0.1, 0.4, 0.5]
-    cluster_coeff = 0.8
-    time_horizon = 10
     transformation = exp
 
-    data = EpiData(gen_int, delay_int, cluster_coeff, time_horizon, transformation)
-    rt_model = ExpGrowthRate(data)
+    data = EpiData(gen_int, transformation)
+    log_init_incidence_prior = Normal()
+    rt_model = ExpGrowthRate(data, log_init_incidence_prior)
 
+    #Example incidence data
     recent_incidence = [10.0, 20.0, 30.0]
     log_init = log(5.0)
     rt = [log(recent_incidence[1]) - log_init; diff(log.(recent_incidence))]
 
-    @test rt_model(rt, log_init) ≈ recent_incidence
+    #Check log_init is sampled from the correct distribution
+    sample_init_inc = sample(EpiAware.generate_latent_infs(rt_model, rt), Prior(), 1000) |>
+                      chn -> chn[:init_incidence] |>
+                             Array |>
+                             vec
+
+    ks_test_pval = ExactOneSampleKSTest(sample_init_inc, log_init_incidence_prior) |> pvalue
+    @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented
+
+    #Check that the generated incidence is correct given correct initialisation
+    mdl_incidence = generated_quantities(EpiAware.generate_latent_infs(rt_model, rt),
+        (init_incidence = log_init,))
+    @test mdl_incidence ≈ recent_incidence
 end
 
-@testitem "DirectInfections function" begin
+@testitem "generate_latent_infs dispatched on DirectInfections" begin
+    using Distributions, Turing, HypothesisTests, DynamicPPL
     gen_int = [0.2, 0.3, 0.5]
-    delay_int = [0.1, 0.4, 0.5]
-    cluster_coeff = 0.8
-    time_horizon = 10
     transformation = exp
 
-    data = EpiData(gen_int, delay_int, cluster_coeff, time_horizon, transformation)
-    direct_inf_model = DirectInfections(data)
+    data = EpiData(gen_int, transformation)
+    log_init_incidence_prior = Normal()
 
+    direct_inf_model = DirectInfections(data, log_init_incidence_prior)
+
+    log_init_scale = log(1.0)
     log_incidence = [10, 20, 30] .|> log
+    expected_incidence = exp.(log_init_scale .+ log_incidence)
+
+    #Check log_init is sampled from the correct distribution
+    sample_init_inc = sample(
+        EpiAware.generate_latent_infs(direct_inf_model, log_incidence),
+        Prior(), 1000) |>
+                      chn -> chn[:init_incidence] |>
+                             Array |>
+                             vec
+
+    ks_test_pval = ExactOneSampleKSTest(sample_init_inc, log_init_incidence_prior) |> pvalue
+    @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented
+
+    #Check that the generated incidence is correct given correct initialisation
+    mdl_incidence = generated_quantities(
+        EpiAware.generate_latent_infs(direct_inf_model,
+            log_incidence),
+        (init_incidence = log_init_scale,))
+
+    @test mdl_incidence ≈ expected_incidence
+end
+@testitem "generate_latent_infs function: default" begin
+    latent_process = [0.1, 0.2, 0.3]
+    init_incidence = 10.0
+
+    struct TestEpiModel <: EpiAware.AbstractEpiModel
+    end
+
+    @test isnothing(EpiAware.generate_latent_infs(TestEpiModel(), latent_process))
+end
+@testitem "generate_latent_infs dispatched on Renewal" begin
+    using Distributions, Turing, HypothesisTests, DynamicPPL, LinearAlgebra
+    gen_int = [0.2, 0.3, 0.5]
+    transformation = exp
+
+    data = EpiData(gen_int, transformation)
+    log_init_incidence_prior = Normal()
+
+    renewal_model = Renewal(data, log_init_incidence_prior)
+
+    #Actual Rt
+    Rt = [1.0, 1.2, 1.5, 1.5, 1.5]
+    log_Rt = log.(Rt)
+    initial_incidence = [1.0, 1.0, 1.0]#aligns with initial exp growth rate of 0.
+
+    #Check log_init is sampled from the correct distribution
+    @time sample_init_inc = sample(EpiAware.generate_latent_infs(renewal_model, log_Rt),
+        Prior(), 1000) |>
+                            chn -> chn[:init_incidence] |>
+                                   Array |>
+                                   vec
+
+    ks_test_pval = ExactOneSampleKSTest(sample_init_inc, log_init_incidence_prior) |> pvalue
+    @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented
+
+    #Check that the generated incidence is correct given correct initialisation
+    #Check first three days "by hand"
+    mdl_incidence = generated_quantities(
+        EpiAware.generate_latent_infs(renewal_model,
+            log_Rt), (init_incidence = 0.0,))
 
-    expected_incidence = exp.(log_incidence)
+    day1_incidence = dot(initial_incidence, gen_int) * Rt[1]
+    day2_incidence = dot(initial_incidence, gen_int) * Rt[2]
+    day3_incidence = dot([day2_incidence, 1.0, 1.0], gen_int) * Rt[3]
 
-    @test direct_inf_model(log_incidence, 0.0) ≈ expected_incidence
+    @test mdl_incidence[1:3] ≈ [day1_incidence, day2_incidence, day3_incidence]
 end
diff --git a/EpiAware/test/test_initialisation.jl b/EpiAware/test/test_initialisation.jl
deleted file mode 100644
index 8cc96fec8..000000000
--- a/EpiAware/test/test_initialisation.jl
+++ /dev/null
@@ -1,21 +0,0 @@
-@testitem "Testing default_initialisation_prior" begin
-    using Distributions
-    prior = EpiAware.default_initialisation_prior()
-
-    @test haskey(prior, :I0_prior)
-    @test typeof(prior[:I0_prior]) <: Normal
-end
-
-@testitem "Testing initialize_incidence" begin
-    using Distributions, Turing
-    using HypothesisTests: ExactOneSampleKSTest, pvalue
-    initialisation_prior = (; I0_prior = Normal())
-    I0_model = EpiAware.initialize_incidence(; initialisation_prior...)
-
-    n_samples = 2000
-    I0_samples = [rand(I0_model) for _ in 1:n_samples] .|> x -> x[:_I0]
-    #Check that the samples are drawn from the correct distribution
-    ks_test_pval = ExactOneSampleKSTest(I0_samples, initialisation_prior.I0_prior) |> pvalue
-
-    @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented
-end
diff --git a/EpiAware/test/test_latent-processes.jl b/EpiAware/test/test_latent-processes.jl
index baa8f500c..a07451ccc 100644
--- a/EpiAware/test/test_latent-processes.jl
+++ b/EpiAware/test/test_latent-processes.jl
@@ -5,7 +5,9 @@
 
     n = 5
     priors = EpiAware.default_rw_priors()
-    model = EpiAware.random_walk(n; priors...)
+    rw_process = EpiAware.RandomWalkLatentProcess(Normal(0.0, 1.0),
+        truncated(Normal(0.0, 0.05), 0.0, Inf))
+    model = EpiAware.generate_latent_process(rw_process, n)
     fixed_model = fix(model, (σ²_RW = 1.0, init_rw_value = 0.0)) #Fixing the standard deviation of the random walk process
     n_samples = 1000
     samples_day_5 = sample(fixed_model, Prior(), n_samples) |>
@@ -18,14 +20,21 @@
 end
 @testitem "Testing default_rw_priors" begin
     @testset "var_RW_prior" begin
-        priors = default_rw_priors()
+        priors = EpiAware.default_rw_priors()
         var_RW = rand(priors[:var_RW_prior])
         @test var_RW >= 0.0
     end
 
     @testset "init_rw_value_prior" begin
-        priors = default_rw_priors()
+        priors = EpiAware.default_rw_priors()
         init_rw_value = rand(priors[:init_rw_value_prior])
         @test typeof(init_rw_value) == Float64
     end
 end
+@testset "Testing RandomWalkLatentProcess constructor" begin
+    init_prior = Normal(0.0, 1.0)
+    var_prior = truncated(Normal(0.0, 0.05), 0.0, Inf)
+    rw_process = RandomWalkLatentProcess(init_prior, var_prior)
+    @test rw_process.init_prior == init_prior
+    @test rw_process.var_prior == var_prior
+end
diff --git a/EpiAware/test/test_models.jl b/EpiAware/test/test_models.jl
index 14ea43345..8cb08cdf3 100644
--- a/EpiAware/test/test_models.jl
+++ b/EpiAware/test/test_models.jl
@@ -1,82 +1,116 @@
 
-@testitem "direct infections with RW latent process" begin
+@testitem "`make_epi_inference_model` with direct infections and RW latent process runs" begin
     using Distributions, Turing, DynamicPPL
     # Define test inputs
     y_t = missing # Data will be generated from the model
-    data = EpiData([0.2, 0.3, 0.5], [0.1, 0.4, 0.5], 0.8, 10, exp)
+    data = EpiData([0.2, 0.3, 0.5], exp)
     pos_shift = 1e-6
 
-    epimodel = DirectInfections(data)
-    rwp = random_walk_process()
-    obs_mdl = delay_observations_model()
-    # Call the function
-    test_mdl = make_epi_inference_model(y_t, epimodel, rwp, obs_mdl; pos_shift)
-
-    # Define expected outputs for a conditional model
-    # Underlying log-infections are const value 1 for all time steps and
-    # any other unfixed parameters
-
-    fixed_test_mdl = fix(
-        test_mdl, (init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05))
-    X = rand(fixed_test_mdl)
-    expected_I_t = [1.0 for _ in 1:(epimodel.data.time_horizon)]
-    gen = generated_quantities(fixed_test_mdl, rand(fixed_test_mdl))
-
-    # Perform tests
-    @test gen.I_t ≈ expected_I_t
+    #Define the epimodel
+    epimodel = DirectInfections(data, Normal())
+
+    #Define the latent process model
+    rwp = EpiAware.RandomWalkLatentProcess(Normal(0.0, 1.0),
+        truncated(Normal(0.0, 0.05), 0.0, Inf))
+
+    #Define the observation model
+    delay_distribution = Gamma(2.0, 5 / 2)
+    time_horizon = 365
+    D_delay = 14.0
+    Δd = 1.0
+
+    obs_model = EpiAware.DelayObservations(delay_distribution = delay_distribution,
+        time_horizon = time_horizon,
+        neg_bin_cluster_factor_prior = Gamma(5, 0.05 / 5),
+        D_delay = D_delay,
+        Δd = Δd)
+
+    # Create full epi model and sample from it
+    test_mdl = make_epi_inference_model(y_t, time_horizon; epimodel = epimodel,
+        latent_process_model = rwp,
+        observation_model = obs_model, pos_shift)
+    gen = generated_quantities(test_mdl, rand(test_mdl))
+
+    #Check model sampled
+    @test eltype(gen.generated_y_t) <: Integer
+    @test eltype(gen.I_t) <: AbstractFloat
+    @test length(gen.I_t) == time_horizon
 end
 
-@testitem "exp growth with RW latent process" begin
+@testitem "`make_epi_inference_model` with Exp growth rate and RW latent process runs" begin
     using Distributions, Turing, DynamicPPL
     # Define test inputs
-    y_t = missing # Data will be generated from the model
-    data = EpiData([0.2, 0.3, 0.5], [0.1, 0.4, 0.5], 0.8, 10, exp)
+    y_t = missing# rand(1:10, 365) # Data will be generated from the model
+    data = EpiData([0.2, 0.3, 0.5], exp)
     pos_shift = 1e-6
 
-    epimodel = ExpGrowthRate(data)
-    rwp = random_walk_process()
-    obs_mdl = delay_observations_model()
-
-    # Call the function
-    test_mdl = make_epi_inference_model(y_t, epimodel, rwp, obs_mdl; pos_shift)
-
-    # Define expected outputs for a conditional model
-    # Underlying log-infections are const value 1 for all time steps and
-    # any other unfixed parameters
-
-    fixed_test_mdl = fix(
-        test_mdl, (init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05))
-    X = rand(fixed_test_mdl)
-    expected_I_t = [1.0 for _ in 1:(epimodel.data.time_horizon)]
-    gen = generated_quantities(fixed_test_mdl, rand(fixed_test_mdl))
-
-    # # Perform tests
-    @test gen.I_t ≈ expected_I_t
+    #Define the epimodel
+    epimodel = EpiAware.ExpGrowthRate(data, Normal())
+
+    #Define the latent process model
+    r_3 = log(2) / 3.0
+    rwp = EpiAware.RandomWalkLatentProcess(
+        truncated(Normal(0.0, r_3 / 3), -r_3, r_3), # 3 day doubling time at 3 sigmas in prior
+        truncated(Normal(0.0, 0.01), 0.0, 0.1))
+
+    #Define the observation model - no delay model
+    time_horizon = 5
+    obs_model = EpiAware.DelayObservations([1.0],
+        time_horizon,
+        truncated(Gamma(5, 0.05 / 5), 1e-3, 1.0))
+
+    # Create full epi model and sample from it
+    test_mdl = make_epi_inference_model(y_t,
+        time_horizon;
+        epimodel = epimodel,
+        latent_process_model = rwp,
+        observation_model = obs_model,
+        pos_shift)
+
+    chn = sample(test_mdl, Prior(), 1000)
+    gens = generated_quantities(test_mdl, chn)
+
+    #Check model sampled
+    @test eltype(gens[1].generated_y_t) <: Integer
+    @test eltype(gens[1].I_t) <: AbstractFloat
+    @test length(gens[1].I_t) == time_horizon
 end
 
-@testitem "Renewal with RW latent process" begin
+@testitem "`make_epi_inference_model` with Renewal and RW latent process runs" begin
     using Distributions, Turing, DynamicPPL
     # Define test inputs
-    y_t = missing # Data will be generated from the model
-    data = EpiData([0.2, 0.3, 0.5], [0.1, 0.4, 0.5], 0.8, 10, exp)
+    y_t = missing# rand(1:10, 365) # Data will be generated from the model
+    data = EpiData([0.2, 0.3, 0.5], exp)
     pos_shift = 1e-6
 
-    epimodel = Renewal(data)
-    rwp = random_walk_process()
-    obs_mdl = delay_observations_model()
-    # Call the function
-    test_mdl = make_epi_inference_model(y_t, epimodel, rwp, obs_mdl; pos_shift)
-
-    # Define expected outputs for a conditional model
-    # Underlying log-infections are const value 1 for all time steps and
-    # any other unfixed parameters
-
-    fixed_test_mdl = fix(
-        test_mdl, (init = log(1.0), σ²_RW = 0.0, neg_bin_cluster_factor = 0.05))
-    X = rand(fixed_test_mdl)
-    expected_I_t = [1.0 for _ in 1:(epimodel.data.time_horizon)]
-    gen = generated_quantities(fixed_test_mdl, rand(fixed_test_mdl))
-
-    # # Perform tests
-    @test gen.I_t ≈ expected_I_t
+    #Define the epimodel
+    epimodel = EpiAware.Renewal(data, Normal())
+
+    #Define the latent process model
+    r_3 = log(2) / 3.0
+    rwp = EpiAware.RandomWalkLatentProcess(
+        truncated(Normal(0.0, r_3 / 3), -r_3, r_3), # 3 day doubling time at 3 sigmas in prior
+        truncated(Normal(0.0, 0.01), 0.0, 0.1))
+
+    #Define the observation model - no delay model
+    time_horizon = 5
+    obs_model = EpiAware.DelayObservations([1.0],
+        time_horizon,
+        truncated(Gamma(5, 0.05 / 5), 1e-3, 1.0))
+
+    # Create full epi model and sample from it
+    test_mdl = make_epi_inference_model(y_t,
+        time_horizon;
+        epimodel = epimodel,
+        latent_process_model = rwp,
+        observation_model = obs_model,
+        pos_shift)
+
+    chn = sample(test_mdl, Prior(), 1000)
+    gens = generated_quantities(test_mdl, chn)
+
+    #Check model sampled
+    @test eltype(gens[1].generated_y_t) <: Integer
+    @test eltype(gens[1].I_t) <: AbstractFloat
+    @test length(gens[1].I_t) == time_horizon
 end
diff --git a/EpiAware/test/test_observation-processes.jl b/EpiAware/test/test_observation-processes.jl
index 3aac3d3ea..d4b2ee7e8 100644
--- a/EpiAware/test/test_observation-processes.jl
+++ b/EpiAware/test/test_observation-processes.jl
@@ -4,25 +4,21 @@
 
     # Set up test data with fixed infection
     I_t = [10.0, 20.0, 30.0]
+    obs_prior = EpiAware.default_delay_obs_priors()
 
-    # Replace with your own implementation of AbstractEpiModel
     # Delay kernel is just event observed on same day
-    data = EpiData([0.2, 0.3, 0.5], [1.0], 0.8, 3, exp)
-    epimodel = DirectInfections(data)
+    delay_obs = EpiAware.DelayObservations([1.0], length(I_t),
+        obs_prior[:neg_bin_cluster_factor_prior])
 
     # Set up priors
-    priors = default_delay_obs_priors()
     neg_bin_cf = 0.05
 
     # Call the function
-    mdl = EpiAware.delay_observations(
+    mdl = EpiAware.generate_observations(delay_obs,
         missing,
-        I_t,
-        epimodel;
-        pos_shift = 1e-6,
-        priors...
-    )
-    fix_mdl = fix(mdl, neg_bin_cluster_factor = neg_bin_cf) # Effectively Poisson sampling
+        I_t;
+        pos_shift = 1e-6)
+    fix_mdl = fix(mdl, (neg_bin_cluster_factor = neg_bin_cf,))
 
     n_samples = 2000
     first_obs = sample(fix_mdl, Prior(), n_samples) |>
@@ -41,3 +37,45 @@
     var_pval = VarianceFTest(first_obs, direct_samples) |> pvalue
     @test var_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented
 end
+
+@testitem "Testing DelayObservations struct" begin
+    using Distributions
+
+    # Test case 1
+    delay_int = [0.2, 0.3, 0.5]
+    time_horizon = 30
+    obs_prior = EpiAware.default_delay_obs_priors()
+
+    obs_model = EpiAware.DelayObservations(delay_int, time_horizon,
+        obs_prior[:neg_bin_cluster_factor_prior])
+
+    @test size(obs_model.delay_kernel) == (time_horizon, time_horizon)
+    @test obs_model.neg_bin_cluster_factor_prior == obs_prior[:neg_bin_cluster_factor_prior]
+
+    # Test case 2
+    delay_distribution = Uniform(0.0, 20.0)
+    time_horizon = 365
+    D_delay = 10.0
+    Δd = 1.0
+
+    obs_model = EpiAware.DelayObservations(delay_distribution = delay_distribution,
+        time_horizon = time_horizon,
+        neg_bin_cluster_factor_prior = obs_prior[:neg_bin_cluster_factor_prior],
+        D_delay = D_delay,
+        Δd = Δd)
+
+    @test size(obs_model.delay_kernel) == (time_horizon, time_horizon)
+    @test obs_model.neg_bin_cluster_factor_prior == obs_prior[:neg_bin_cluster_factor_prior]
+end
+
+@testitem "Testing generate_observations default" begin
+    struct TestObsModel <: EpiAware.AbstractObservationModel
+    end
+
+    @test try
+        EpiAware.generate_observations(TestObsModel(), missing, missing; pos_shift = 1e-6)
+        true
+    catch
+        false
+    end
+end
diff --git a/EpiAware/test/test_utilities.jl b/EpiAware/test/test_utilities.jl
index 47e869ca6..8ae9b58d2 100644
--- a/EpiAware/test/test_utilities.jl
+++ b/EpiAware/test/test_utilities.jl
@@ -51,13 +51,11 @@ end
     @testset "Test case 4" begin
         dist = Exponential(1.0)
         expected_pmf = [(exp(-(t - 1)) - exp(-t)) / (1 - exp(-5)) for t in 1:5]
-        pmf = create_discrete_pmf(
-            dist,
+        pmf = create_discrete_pmf(dist,
             Val(:single_censored);
             primary_approximation_point = 0.0,
             Δd = 1.0,
-            D = 5.0
-        )
+            D = 5.0)
         @test pmf≈expected_pmf atol=1e-15
     end
 
@@ -101,13 +99,11 @@ end
     @testset "Test case 1" begin
         delay_int = [0.2, 0.5, 0.3]
         time_horizon = 5
-        expected_K = SparseMatrixCSC(
-            [0.2 0 0 0 0
-             0.5 0.2 0 0 0
-             0.3 0.5 0.2 0 0
-             0 0.3 0.5 0.2 0
-             0 0 0.3 0.5 0.2],
-        )
+        expected_K = SparseMatrixCSC([0.2 0 0 0 0
+                                      0.5 0.2 0 0 0
+                                      0.3 0.5 0.2 0 0
+                                      0 0.3 0.5 0.2 0
+                                      0 0 0.3 0.5 0.2])
         K = EpiAware.generate_observation_kernel(delay_int, time_horizon)
         @test K == expected_K
     end
@@ -150,12 +146,10 @@ end
     @testset "Test case 2" begin
         r = 0
         w = [0.1, 0.2, 0.3, 0.4]
-        expected_result = -(
-            0.1 * 1 * exp(-0 * 1) +
-            0.2 * 2 * exp(-0 * 2) +
-            0.3 * 3 * exp(-0 * 3) +
-            0.4 * 4 * exp(-0 * 4)
-        )
+        expected_result = -(0.1 * 1 * exp(-0 * 1) +
+                            0.2 * 2 * exp(-0 * 2) +
+                            0.3 * 3 * exp(-0 * 3) +
+                            0.4 * 4 * exp(-0 * 4))
         result = EpiAware.dneg_MGF_dr(r, w)
         @test result≈expected_result atol=1e-15
     end