Merge pull request #114 from OHDSI/mockCohort

fixes recoredPerson = 1 behaviour

R/mockCohort.R

@@ -51,7 +51,7 @@ mockCohort <- function(cdm,
                        numberCohorts = 1,
                        cohortName = paste0("cohort_", seq_len(numberCohorts)),
                        recordPerson = 1,
-                       seed = 1) {
+                       seed = NULL) {
   # initial checks
   checkInput(
     cdm = cdm,
@@ -85,7 +85,7 @@ mockCohort <- function(cdm,
   numberRows <-
     recordPerson * (cdm$person |> dplyr::tally() |> dplyr::pull()) |> round()
 
-  numberRows <- numberRows * 1.2
+  numberRows <- (numberRows * 1.2) |> round()
   rows_to_keep <- sum(numberRows / 1.2)
 
 
@@ -146,18 +146,29 @@ mockCohort <- function(cdm,
     dplyr::ungroup() |>
     dplyr::select(-"next_observation") |>
     stats::na.omit() |>
-    dplyr::distinct() |>
-    dplyr::slice(1:rows_to_keep)
-
-
-
+    dplyr::distinct()
+
+#correct cohort count
+  if(nrow(cohort) > 0) {
+    cohort_id <- cohort |>
+      dplyr::pull("cohort_definition_id") |>
+      unique() |> as.integer()
+
+    numberRows <- (numberRows / 1.2) |> round()
+
+    cohort <- purrr::map(
+      cohort_id,
+      \(x) cohort |>
+        dplyr::filter(.data$cohort_definition_id == x) |>
+        dplyr::slice(1:numberRows[x])
+    ) |> dplyr::bind_rows()
+  }
   # generate cohort set table
 
   cohortName <- snakecase::to_snake_case(cohortName)
 
   cohortSetTable <- dplyr::tibble(cohort_definition_id = cohortId,
                                   cohort_name = cohortName)
-
   # create class
 
   cdm <-

tests/testthat/test-mockCohort.R

@@ -18,7 +18,7 @@ test_that("mock cohort simple test", {
     omock::mockObservationPeriod() |>
     omock::mockCohort(recordPerson = 2)
 
-  testthat::expect_true(cdm$cohort |> dplyr::tally() == 200)
+  expect_true(cdm$cohort |> dplyr::tally() == 200)
 
   expect_no_error(mockCdmReference() |>
     mockPerson(nPerson = 100) |>
@@ -35,3 +35,21 @@ test_that("mock cohort simple test", {
 
 
 })
+
+test_that("cohort count", {
+  cdm <- omock::emptyCdmReference(cdmName = "mock") |>
+    omock::mockPerson(nPerson = 100) |>
+    omock::mockObservationPeriod() |>
+    omock::mockCohort(recordPerson = 1,
+                      numberCohorts = 3,
+                      seed = 1)
+
+  expect_true(all(
+    cdm$cohort |>
+      dplyr::group_by(cohort_definition_id) |>
+      dplyr::tally() |>
+      dplyr::pull(n) == c(100, 100, 100)
+  ))
+
+})
+