diff --git a/R/plotting.R b/R/plotting.R
index 71f7a0d..57a2083 100644
--- a/R/plotting.R
+++ b/R/plotting.R
@@ -144,13 +144,13 @@ plot.disag_model <- function(x, include_iid = FALSE, ...){
       geom_abline(intercept = 0, slope = 1, color = 'blue') +
       scale_color_manual(values = c("Without IID" = "black")) +
       ggtitle(title) +
-      labs(color = NULL) +
+      labs(x = "Observed", y = "Predicted", color = NULL) +
       theme(legend.position.inside = c(0, 1),
             legend.justification = c(0, 1))
     if (include_iid){
       obspred$scales$scales <- list()
       obspred <- obspred +
-        geom_point(data, aes(x = obs, y = pred, color = "With IID")) +
+        geom_point(data = data, aes(x = obs, y = pred, color = "With IID")) +
         scale_color_manual(values = c("Without IID" = "black", "With IID" = "red"))
     }
   } else {
diff --git a/tests/testthat/test-plotting.R b/tests/testthat/test-plotting.R
index 3240be3..78b554e 100644
--- a/tests/testthat/test-plotting.R
+++ b/tests/testthat/test-plotting.R
@@ -45,12 +45,16 @@ test_that("Check plot.disag_model function works as expected", {
 
   p2 <- plot(fit_result_nofield)
 
+  p3 <- plot(result, include_iid = TRUE)
+
   expect_is(p1, 'list')
   expect_equal(length(p1), 2)
 
   expect_is(p2, 'list')
   expect_equal(length(p2), 2)
 
+  expect_is(p3, 'list')
+  expect_equal(length(p3), 2)
 
 })