diff --git a/DESCRIPTION b/DESCRIPTION
index 3a59c3b..a337e02 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,7 +1,7 @@
 Package: grstat
 Type: Package
 Title: Clinical Research Tools
-Version: 0.1.0.9005
+Version: 0.1.0.9006
 Authors@R: c(
       person("Dan", "Chaltiel", role = c("aut", "cre"),
              email = "dan.chaltiel@gustaveroussy.fr",
diff --git a/R/ae_table_soc.R b/R/ae_table_soc.R
index 177a9bc..1acd811 100644
--- a/R/ae_table_soc.R
+++ b/R/ae_table_soc.R
@@ -294,30 +294,33 @@ butterfly_plot = function(
     arm="ARM", subjid="SUBJID", soc="AESOC"
 ){
   check_dots_empty()
-  sort_by = arg_match(sort_by)
-
+  assert_not_null(df_ae, df_enrol, sort_by, subjid, soc)
   assert_names_exists(df_ae, lst(subjid, soc, severe))
   assert_names_exists(df_enrol, lst(subjid, arm))
+  sort_by = arg_match(sort_by)
 
   df_ae = df_ae %>%
     select(subjid_=any_of2(subjid), soc_=any_of2(soc),
            severe_=any_of2(severe)) %>%
     mutate(severe_ = if(is.null(severe)) NA else severe_)
   df_enrol = df_enrol %>%
+    mutate(across(any_of2(arm), factor)) %>%
     select(subjid_=any_of2(subjid), arm_=any_of2(arm))
+
+  arms = df_enrol[["arm_"]]
+  n_arms = n_distinct(arms, na.rm=TRUE)
+  if(n_arms!=2){
+    if(is.null(arms)) arms = "NULL"
+    cli_abort(c("{.fn grstat::butterfly_plot} needs exactly 2 arms.",
+                i="Found {n_arms} arm{?s} in column {.arg {arm}}: {.val {unique(arms)}}"),
+              class="edc_butterfly_two_arms_error")
+  }
+
   df = df_enrol %>%
     left_join(df_ae, by="subjid_") %>%
     filter(!is.na(soc_))  %>%
     arrange(subjid_)
 
-  if(!is.factor(df_enrol$arm_)) df_enrol$arm_ = factor(df_enrol$arm_)
-
-  arms = df_enrol$arm_ %>% unique() %>% na.omit()
-  if(length(arms)!=2){
-    cli_abort(c("{.fn EDCimport::butterfly_plot} needs exactly 2 arms.",
-                i="Arms: {.val {arms}}"),
-              class="edc_butterfly_two_arms_error")
-  }
   if(!is.null(severe)){
     if(!is.logical(df_ae$severe_)){
       cli_abort(c("{.arg severe} should be a logical column, not a {.type {df_ae$severe_}}. Did you forget to mutate it with `==`?"),
diff --git a/R/assertions.R b/R/assertions.R
index 95cc157..af8deb1 100644
--- a/R/assertions.R
+++ b/R/assertions.R
@@ -51,6 +51,14 @@ assert_names_exists = function(df, l){
   }
 }
 
+assert_not_null = function(...){
+  nulls = lst(...) %>% keep(is.null) %>% names()
+  if(length(nulls)>0){
+    cli_abort("Variable{?s} {.arg {nulls}} cannot be NULL.",
+              class="grstat_var_null")
+  }
+}
+
 
 # Misc ----------------------------------------------------------------------------------------
 
diff --git a/tests/testthat/test-ae-plots.R b/tests/testthat/test-ae-plots.R
index e2ce731..0aa28c5 100644
--- a/tests/testthat/test-ae-plots.R
+++ b/tests/testthat/test-ae-plots.R
@@ -1,7 +1,7 @@
 
 test_that("ae_plot_grade() works", {
   tm = grstat_example()
-  attach(tm)
+  attach(tm, warn.conflicts=FALSE)
 
   p = ae_plot_grade(df_ae=ae, df_enrol=enrolres)
   vdiffr::expect_doppelganger("ae-plot-grade-1", p)
@@ -19,7 +19,7 @@ test_that("ae_plot_grade() works", {
 
 test_that("ae_plot_grade_sum() works", {
   tm = grstat_example()
-  attach(tm)
+  attach(tm, warn.conflicts=FALSE)
 
   p = ae_plot_grade_sum(df_ae=ae, df_enrol=enrolres)
   vdiffr::expect_doppelganger("ae-plot-grade-sum-1", p)
@@ -35,10 +35,10 @@ test_that("ae_plot_grade_sum() works", {
 
 test_that("butterfly_plot() works", {
   tm = grstat_example()
-  attach(tm)
+  attach(tm, warn.conflicts=FALSE)
   ae2 = ae %>%
     mutate(serious = sae=="Yes",
-           bad_serious = sae=="foobar",)
+           bad_serious = sae=="foobar")
 
   p = butterfly_plot(ae2, df_enrol=enrolres)
   vdiffr::expect_doppelganger("butterfly-plot-1", p)
@@ -46,20 +46,40 @@ test_that("butterfly_plot() works", {
   vdiffr::expect_doppelganger("butterfly-plot-2", p)
   p = butterfly_plot(ae2, df_enrol=enrolres, range_min=1)
   vdiffr::expect_doppelganger("butterfly-plot-3", p)
+})
+
+test_that("butterfly_plot() errors", {
+  tm = grstat_example()
+  attach(tm, warn.conflicts=FALSE)
+  ae2 = ae %>%
+    mutate(serious = sae=="Yes",
+           bad_serious = sae=="foobar")
 
+  # Warnings ------------------------------------------------------
 
-  # Warnings
+  # there is no SAE (probably an error)
   ae2 %>%
     butterfly_plot(df_enrol=enrolres, severe="bad_serious") %>%
     expect_warning(class="edc_butterfly_serious_false_warning")
 
-  # Errors
+  # Errors   ------------------------------------------------------
+
+  # `severe` is not logical
   ae %>%
     butterfly_plot(df_enrol=enrolres, severe="sae") %>%
     expect_error(class="edc_butterfly_serious_lgl_error")
 
+  # not exactly 2 arms
   ae %>%
     butterfly_plot(df_enrol=enrolres, arm="crfname") %>%
     expect_error(class="edc_butterfly_two_arms_error")
 
+  # NULL arguments
+  ae %>%
+    butterfly_plot(df_enrol=NULL, sort_by=NULL, subjid=NULL, soc=NULL) %>%
+    expect_error(class="grstat_var_null")
+  ae %>%
+    butterfly_plot(df_enrol=enrolres, arm=NULL) %>%
+    expect_error(class="edc_butterfly_two_arms_error")
+
 })