diff --git a/R/pred_pred_f.R b/R/pred_pred_f.R
index defa441..650d021 100644
--- a/R/pred_pred_f.R
+++ b/R/pred_pred_f.R
@@ -45,6 +45,9 @@ pred_pred <- function(fit, aggregation, cases, predict_iid, uncertain = FALSE, N
 
   if (uncertain){
     prediction$uncertainty <- disaggregation::predict_uncertainty(fit, predict_iid = predict_iid, N = N, CI = CI)
+    prediction$uncertainty_lower <- prediction$uncertainty$predictions_ci$`lower CI`
+    prediction$uncertainty_upper <- prediction$uncertainty$predictions_ci$`upper CI`
+    prediction$uncertainty <- NULL
   }
 
   names(prediction)[which(names(prediction) == "prediction")] <- "prediction (rate)"
@@ -67,8 +70,8 @@ pred_pred <- function(fit, aggregation, cases, predict_iid, uncertain = FALSE, N
       prediction$iid <- terra::wrap(prediction$iid)
     }
     if (uncertain){
-      prediction$uncertainty_lower <- terra::wrap(prediction$uncertainty$predictions_ci$`lower CI`)
-      prediction$uncertainty_upper <- terra::wrap(prediction$uncertainty$predictions_ci$`upper CI`)
+      prediction$uncertainty_lower <- terra::wrap(prediction$uncertainty_lower)
+      prediction$uncertainty_upper <- terra::wrap(prediction$uncertainty_upper)
     }
   }
 
diff --git a/inst/shiny/modules/core_load.R b/inst/shiny/modules/core_load.R
index d3a52e8..5feb62a 100644
--- a/inst/shiny/modules/core_load.R
+++ b/inst/shiny/modules/core_load.R
@@ -63,18 +63,19 @@ core_load_module_server <- function(id, common, modules, map, COMPONENT_MODULES,
       common$agg_prep <- unwrap_terra(common$agg_prep)
       common$agg_prep_lores <- unwrap_terra(common$agg_prep_lores)
       common$prep$covariate_rasters <- unwrap_terra(common$prep$covariate_rasters)
-      common$pred$field <- unwrap_terra(common$pred$field)
+      common$fit$data$covariate_rasters <- unwrap_terra(common$fit$data$covariate_rasters)
       common$pred$`prediction (rate)` <- unwrap_terra(common$pred$`prediction (rate)`)
       common$pred$`prediction (cases)` <- unwrap_terra(common$pred$`prediction (cases)`)
+      common$pred$covariates <- unwrap_terra(common$pred$covariates)
       common$pred$iid <- unwrap_terra(common$pred$iid)
-      common$fit$data$covariate_rasters <- unwrap_terra(common$fit$data$covariate_rasters)
+      common$pred$field <- unwrap_terra(common$pred$field)
+      common$pred$uncertainty_lower <- unwrap_terra(common$pred$uncertainty_lower)
+      common$pred$uncertainty_upper <- unwrap_terra(common$pred$uncertainty_upper)
       common$transfer$agg <- unwrap_terra(common$transfer$agg)
       common$transfer$cases <- unwrap_terra(common$transfer$cases)
       common$transfer$prediction <- unwrap_terra(common$transfer$prediction)
       common$transfer$field <- unwrap_terra(common$transfer$field)
       common$transfer$covariates <- unwrap_terra(common$transfer$covariates)
-      common$pred$uncertainty_lower <- unwrap_terra(common$pred$uncertainty_lower)
-      common$pred$uncertainty_upper <- unwrap_terra(common$pred$uncertainty_upper)
 
       #restore map and results for used modules
       for (used_module in names(common$meta)){
@@ -106,7 +107,9 @@ core_load_module_server <- function(id, common, modules, map, COMPONENT_MODULES,
       }
 
       close_loading_modal()
-      common$logger |> writeLog(type = "info", "The previous session has been loaded successfully")
+      if (isFALSE(getOption("shiny.testmode"))) {
+        common$logger |> writeLog(type = "info", "The previous session has been loaded successfully")
+      }
     })
 
   }
diff --git a/inst/shiny/modules/core_save.R b/inst/shiny/modules/core_save.R
index d26ed04..5dff84c 100644
--- a/inst/shiny/modules/core_save.R
+++ b/inst/shiny/modules/core_save.R
@@ -56,19 +56,20 @@ core_save_module_server <- function(id, common, modules, COMPONENTS, main_input)
         common$agg_prep <- wrap_terra(common$agg_prep)
         common$agg_prep_lores <- wrap_terra(common$agg_prep_lores)
         common$prep$covariate_rasters <- wrap_terra(common$prep$covariate_rasters)
-        common$pred$field <- wrap_terra(common$pred$field)
+        common$fit$data$covariate_rasters <- wrap_terra(common$fit$data$covariate_rasters)
         common$pred$`prediction (rate)` <- wrap_terra(common$pred$`prediction (rate)`)
         common$pred$`prediction (cases)` <- wrap_terra(common$pred$`prediction (cases)`)
         common$pred$covariates <- wrap_terra(common$pred$covariates)
         common$pred$iid <- wrap_terra(common$pred$iid)
-        common$fit$data$covariate_rasters <- wrap_terra(common$fit$data$covariate_rasters)
+        common$pred$field <- wrap_terra(common$pred$field)
+        common$pred$uncertainty_lower <- wrap_terra(common$pred$uncertainty_lower)
+        common$pred$uncertainty_upper <- wrap_terra(common$pred$uncertainty_upper)
         common$transfer$agg <- wrap_terra(common$transfer$agg)
         common$transfer$cases <- wrap_terra(common$transfer$cases)
         common$transfer$prediction <- wrap_terra(common$transfer$prediction)
         common$transfer$field <- wrap_terra(common$transfer$field)
         common$transfer$covariates <- wrap_terra(common$transfer$covariates)
-        common$pred$uncertainty_lower <- wrap_terra(common$pred$uncertainty_lower)
-        common$pred$uncertainty_upper <- wrap_terra(common$pred$uncertainty_upper)
+
 
         #save file
         saveRDS(common, file)
@@ -81,19 +82,20 @@ core_save_module_server <- function(id, common, modules, COMPONENTS, main_input)
         common$agg_prep <- unwrap_terra(common$agg_prep)
         common$agg_prep_lores <- unwrap_terra(common$agg_prep_lores)
         common$prep$covariate_rasters <- unwrap_terra(common$prep$covariate_rasters)
-        common$pred$field <- unwrap_terra(common$pred$field)
+        common$fit$data$covariate_rasters <- unwrap_terra(common$fit$data$covariate_rasters)
         common$pred$`prediction (rate)` <- unwrap_terra(common$pred$`prediction (rate)`)
         common$pred$`prediction (cases)` <- unwrap_terra(common$pred$`prediction (cases)`)
         common$pred$covariates <- unwrap_terra(common$pred$covariates)
         common$pred$iid <- unwrap_terra(common$pred$iid)
-        common$fit$data$covariate_rasters <- unwrap_terra(common$fit$data$covariate_rasters)
+        common$pred$field <- unwrap_terra(common$pred$field)
+        common$pred$uncertainty_lower <- unwrap_terra(common$pred$uncertainty_lower)
+        common$pred$uncertainty_upper <- unwrap_terra(common$pred$uncertainty_upper)
         common$transfer$agg <- unwrap_terra(common$transfer$agg)
         common$transfer$cases <- unwrap_terra(common$transfer$cases)
         common$transfer$prediction <- unwrap_terra(common$transfer$prediction)
         common$transfer$field <- unwrap_terra(common$transfer$field)
         common$transfer$covariates <- unwrap_terra(common$transfer$covariates)
-        common$pred$uncertainty_lower <- unwrap_terra(common$pred$uncertainty_lower)
-        common$pred$uncertainty_upper <- unwrap_terra(common$pred$uncertainty_upper)
+
 
         close_loading_modal()
       }
diff --git a/tests/testthat/helper_data.R b/tests/testthat/helper_data.R
index d5a4381..3ce223c 100644
--- a/tests/testthat/helper_data.R
+++ b/tests/testthat/helper_data.R
@@ -42,3 +42,59 @@ if (is_ci){
 } else {
   save_path <- "~/temprds/saved_file.rds"
 }
+
+
+polygons <- list()
+n_polygon_per_side <- 10
+n_polygons <- n_polygon_per_side * n_polygon_per_side
+n_pixels_per_side <- n_polygon_per_side * 2
+
+for(i in 1:n_polygons) {
+  row <- ceiling(i/n_polygon_per_side)
+  col <- ifelse(i %% n_polygon_per_side != 0, i %% n_polygon_per_side, n_polygon_per_side)
+  xmin = 2*(col - 1); xmax = 2*col; ymin = 2*(row - 1); ymax = 2*row
+  polygons[[i]] <- list(cbind(c(xmin, xmax, xmax, xmin, xmin),
+                              c(ymax, ymax, ymin, ymin, ymax)))
+}
+
+polys <- lapply(polygons, sf::st_polygon)
+response_df <- data.frame(area_id = 1:n_polygons, response = runif(n_polygons, min = 0, max = 1000))
+spdf <- sf::st_sf(response_df, geometry = polys)
+
+# Create raster stack
+r <- terra::rast(ncol=n_pixels_per_side, nrow=n_pixels_per_side)
+terra::ext(r) <- terra::ext(spdf)
+r[] <- sapply(1:terra::ncell(r), function(x) rnorm(1, ifelse(x %% n_pixels_per_side != 0, x %% n_pixels_per_side, n_pixels_per_side), 3))
+r2 <- terra::rast(ncol=n_pixels_per_side, nrow=n_pixels_per_side)
+terra::ext(r2) <- terra::ext(spdf)
+r2[] <- sapply(1:terra::ncell(r), function(x) rnorm(1, ceiling(x/n_pixels_per_side), 3))
+cov_stack <- c(r, r2)
+names(cov_stack) <- c('layer1', 'layer2')
+
+test_data <- disaggregation::prepare_data(polygon_shapefile = spdf,
+                          covariate_rasters = cov_stack)
+
+result <- disaggregation::disag_model(test_data,
+                      field = TRUE,
+                      iid = TRUE,
+                      iterations = 100,
+                      family = "poisson",
+                      link = "log")
+
+test_common <- list()
+test_common$shape <- spdf
+test_common$covs_prep <- cov_stack
+test_common$agg_prep <- wrap_terra(r)
+test_common$prep <- test_data
+test_common$covs_prep <- wrap_terra(test_common$covs_prep)
+test_common$prep$covariate_rasters <- wrap_terra(test_common$prep$covariate_rasters)
+test_common$state$main$version = as.character(packageVersion("disagapp"))
+class(test_common) <- "common"
+if (is_ci){
+  test_common_path <- tempfile(fileext = ".rds")
+} else {
+  test_common_path <- "~/temprds/test_common.rds"
+}
+saveRDS(test_common, test_common_path)
+
+
diff --git a/tests/testthat/test-fit_fit.R b/tests/testthat/test-fit_fit.R
new file mode 100644
index 0000000..0a65b8d
--- /dev/null
+++ b/tests/testthat/test-fit_fit.R
@@ -0,0 +1,46 @@
+test_that("Check fit_fit function works as expected", {
+
+  result <- fit_fit(data = test_data,
+                    priors = NULL,
+                    family = "poisson",
+                    link = "log",
+                    iterations = 100,
+                    field = TRUE,
+                    iid = TRUE,
+                    async = FALSE)
+
+  expect_is(result, "disag_model")
+  expect_is(result$data$covariate_rasters, "SpatRaster")
+
+  test_data$covariate_rasters <- terra::wrap(test_data$covariate_rasters)
+
+  result <- fit_fit(data = test_data,
+                    priors = NULL,
+                    family = "poisson",
+                    link = "log",
+                    iterations = 100,
+                    field = TRUE,
+                    iid = TRUE,
+                    async = TRUE)
+
+  expect_is(result, "disag_model")
+  expect_is(result$data$covariate_rasters, "PackedSpatRaster")
+
+})
+
+test_that("Check fit_fit function works in the app", {
+  app <- shinytest2::AppDriver$new(app_dir = system.file("shiny", package = "disagapp"), name = "e2e_fit_fit", timeout = 60000)
+  app$set_inputs(tabs = "intro")
+  app$set_inputs(introTabs = "Load Prior Session")
+  app$upload_file(`core_load-load_session` = test_common_path)
+  app$click("core_load-goLoad_session")
+  app$set_inputs(tabs = "fit")
+  app$set_inputs(fitSel = "fit_fit")
+  app$click(selector = "#fit_fit-run")
+  app$wait_for_value(input = "fit_fit-complete")
+
+  app$set_inputs(main = "Save")
+  app$get_download("core_save-save_session", filename = save_path)
+  common <- readRDS(save_path)
+  expect_is(common$fit, "disag_model")
+})
diff --git a/tests/testthat/test-pred_pred.R b/tests/testthat/test-pred_pred.R
new file mode 100644
index 0000000..d6e21ca
--- /dev/null
+++ b/tests/testthat/test-pred_pred.R
@@ -0,0 +1,101 @@
+# test_that("Check pred_pred function works as expected", {
+#
+#   prediction <- pred_pred(fit = result,
+#                           aggregation = r,
+#                           cases = FALSE,
+#                           predict_iid = FALSE,
+#                           uncertain = FALSE,
+#                           async = FALSE)
+#
+#   expect_is(prediction, "list")
+#   expect_length(prediction, 4)
+#   expect_is(prediction$`prediction (rate)`, "SpatRaster")
+#   expect_null(prediction$iid)
+#   expect_is(prediction$field, "SpatRaster")
+#   expect_is(prediction$covariates, "SpatRaster")
+#
+#   prediction <- pred_pred(fit = result,
+#                           aggregation = r,
+#                           cases = TRUE,
+#                           predict_iid = TRUE,
+#                           uncertain = TRUE,
+#                           N = 100,
+#                           CI = 0.95,
+#                           async = FALSE)
+#
+#   expect_is(prediction, "list")
+#   expect_length(prediction, 7)
+#   expect_is(prediction$`prediction (rate)`, "SpatRaster")
+#   expect_is(prediction$`prediction (cases)`, "SpatRaster")
+#   expect_is(prediction$field, "SpatRaster")
+#   expect_is(prediction$iid, "SpatRaster")
+#   expect_is(prediction$uncertainty_lower, "SpatRaster")
+#   expect_is(prediction$uncertainty_upper, "SpatRaster")
+#   expect_is(prediction$covariates, "SpatRaster")
+#
+#   prediction <- pred_pred(fit = result,
+#                           aggregation = r,
+#                           cases = TRUE,
+#                           predict_iid = TRUE,
+#                           uncertain = TRUE,
+#                           N = 100,
+#                           CI = 0.95,
+#                           async = TRUE)
+#
+#   expect_is(prediction, "list")
+#   expect_length(prediction, 7)
+#   expect_is(prediction$`prediction (rate)`, "PackedSpatRaster")
+#   expect_is(prediction$`prediction (cases)`, "PackedSpatRaster")
+#   expect_is(prediction$field, "PackedSpatRaster")
+#   expect_is(prediction$iid, "PackedSpatRaster")
+#   expect_is(prediction$uncertainty_lower, "PackedSpatRaster")
+#   expect_is(prediction$uncertainty_upper, "PackedSpatRaster")
+#   expect_is(prediction$covariates, "PackedSpatRaster")
+#
+# })
+
+test_that("Check fit_fit function works in the app", {
+
+  # setup common for test
+  test_common$fit <- result
+  test_common$fit$data$covariate_rasters <- wrap_terra(test_common$fit$data$covariate_rasters)
+  test_common$meta$fit_fit$iid <- TRUE
+  saveRDS(test_common, test_common_path)
+
+  app <- shinytest2::AppDriver$new(app_dir = system.file("shiny", package = "disagapp"), name = "e2e_fit_fit", timeout = 60000)
+  app$set_inputs(tabs = "intro")
+  app$set_inputs(introTabs = "Load Prior Session")
+  app$upload_file(`core_load-load_session` = test_common_path)
+  app$click("core_load-goLoad_session")
+
+  app$set_inputs(tabs = "pred")
+  app$set_inputs(predSel = "pred_pred")
+  app$set_inputs("pred_pred-cases" = TRUE)
+  app$set_inputs("pred_pred-iid" = TRUE)
+  app$set_inputs("pred_pred-uncertain" = TRUE)
+  app$click(selector = "#pred_pred-run")
+  app$wait_for_value(input = "pred_pred-complete")
+  app$set_inputs(main = "Save")
+  app$get_download("core_save-save_session", filename = save_path)
+
+  common <- readRDS(save_path)
+  common$pred$`prediction (rate)` <- unwrap_terra(common$pred$`prediction (rate)`)
+  common$pred$`prediction (cases)` <- unwrap_terra(common$pred$`prediction (cases)`)
+  common$pred$field <- unwrap_terra(common$pred$field)
+  common$pred$iid <- unwrap_terra(common$pred$iid)
+  common$pred$uncertainty_lower <- unwrap_terra(common$pred$uncertainty_lower)
+  common$pred$uncertainty_upper <- unwrap_terra(common$pred$uncertainty_upper)
+  common$pred$covariates <- unwrap_terra(common$pred$covariates)
+
+  expect_is(common$pred$`prediction (rate)`, "SpatRaster")
+  expect_is(common$pred$`prediction (cases)`, "SpatRaster")
+  expect_is(common$pred$field, "SpatRaster")
+  expect_is(common$pred$iid, "SpatRaster")
+  expect_is(common$pred$uncertainty_lower, "SpatRaster")
+  expect_is(common$pred$uncertainty_upper, "SpatRaster")
+  expect_is(common$pred$covariates, "SpatRaster")
+})
+
+
+
+