From c1f32fac55e01abe8f2a3391469be6cc2ae7ac0f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Matthias=20Greni=C3=A9?= <matthias.grenie@ens-lyon.fr>
Date: Fri, 5 Apr 2024 10:58:56 +0200
Subject: [PATCH] Fix #106

---
 R/fb_plot_trait_correlation.R                 | 11 ++++----
 .../testthat/test-fb_plot_trait_correlation.R | 27 ++++++++++++++-----
 2 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/R/fb_plot_trait_correlation.R b/R/fb_plot_trait_correlation.R
index 3ba06b9a..b612ed63 100644
--- a/R/fb_plot_trait_correlation.R
+++ b/R/fb_plot_trait_correlation.R
@@ -30,17 +30,18 @@ fb_plot_trait_correlation <- function(
   check_species_categories(species_categories)
   
   # Keep only numeric columns
-  trait_type <- vapply(species_traits[, -1, drop = FALSE], typeof, character(1))
-  
-  numerical_traits <- trait_type %in% c("double", "integer")
+  numerical_traits <- vapply(
+    species_traits[, -1, drop = FALSE], is.numeric, logical(1)
+  )
   
   if (!any(numerical_traits)) {
-    stop("No numerical traits found, cannot plot trait correlations")
+    stop("No numerical traits found, cannot plot trait correlations",
+         call. = FALSE)
   }
   
   if (!all(numerical_traits)) {
     message("Non-numerical traits found, only keeping numerical traits ",
-            "to show trait correlation")
+            "to display trait correlations")
   }
   
   # Subset Traits --------------------------------------------------------------
diff --git a/tests/testthat/test-fb_plot_trait_correlation.R b/tests/testthat/test-fb_plot_trait_correlation.R
index 5763a2b1..e1ab01d0 100644
--- a/tests/testthat/test-fb_plot_trait_correlation.R
+++ b/tests/testthat/test-fb_plot_trait_correlation.R
@@ -1,20 +1,32 @@
 # Initial data -----------------------------------------------------------------
 
-sp_trait <- data.frame(species = letters[1:3], trait1 = letters[1:3],
-                       trait2  = 1:3, trait3 = 3:1)
+sp_trait <- data.frame(
+  species = letters[1:3],
+  trait1  = letters[1:3],
+  trait2  = 1:3,
+  trait3  = 3:1,
+  trait4  = factor(letters[1:3])
+)
 
 
 # Actual Tests -----------------------------------------------------------------
 
 test_that("fb_plot_trait_correlation() works", {
   
-  # Wrong input
+  # Wrong input, only non-numerical traits
   expect_error(
     fb_plot_trait_correlation(sp_trait[, 1:2]),
     "No numerical traits found, cannot plot trait correlations",
     fixed = TRUE
   )
   
+  # Wrong input, only non-numerical traits
+  expect_error(
+    fb_plot_trait_correlation(sp_trait[, c(1:2, 5)]),
+    "No numerical traits found, cannot plot trait correlations",
+    fixed = TRUE
+  )
+  
   # Good input
   expect_silent(res <- fb_plot_trait_correlation(species_traits))
   
@@ -25,7 +37,7 @@ test_that("fb_plot_trait_correlation() works", {
     res <- fb_plot_trait_correlation(sp_trait),
     paste0(
       "Non-numerical traits found, only keeping numerical traits ",
-      "to show trait correlation"
+      "to display trait correlations"
     ),
     fixed = TRUE
   )
@@ -36,14 +48,15 @@ test_that("fb_plot_trait_correlation() works", {
   # Single category 
   expect_silent(
     given_plot <- fb_plot_trait_correlation(
-      sp_trait[, -2], data.frame(species  = sp_trait$species, category = "A")
+      sp_trait[, -c(2, 5)],
+      data.frame(species = sp_trait$species, category = "A")
     )
   )
   
   # Less categories than species
   expect_silent(
     given_plot <- fb_plot_trait_correlation(
-      sp_trait[, -2],
+      sp_trait[, -c(2, 5)],
       data.frame(species  = sp_trait$species, category = c(1, 1, 2))
     )
   )
@@ -51,7 +64,7 @@ test_that("fb_plot_trait_correlation() works", {
   # As many categories as species
   expect_silent(
     given_plot <- fb_plot_trait_correlation(
-      sp_trait[, -2],
+      sp_trait[, -c(2, 5)],
       data.frame(species  = sp_trait$species,
                  category = sp_trait$species)
     )