test_that("test class consistency across cohort operations", {
  dus_cohort <- dplyr::tibble(
    cohort_definition_id = c(1, 1, 1, 2),
    subject_id = c(1, 1, 2, 2),
    cohort_start_date = as.Date(c(
      "2010-04-19", "2020-04-19", "2021-01-14", "2012-05-25"
    )),
    cohort_end_date = as.Date(c(
      "2010-09-19", "2020-08-10", "2021-11-14", "2022-09-25"
    ))
  )
  observation_period <-  dplyr::tibble(observation_period_id = c(1, 2),
                                       person_id = c(1, 2),
                                       observation_period_start_date = c(as.Date("2006-03-11"),
                                                                         as.Date("2006-03-11")),
                                       observation_period_end_date = c(as.Date("2102-04-02"),
                                                                       as.Date("2102-04-02")),
                                       period_type_concept_id = c(0, 0))
  cdm <- mockPatientProfiles(
    connectionDetails,
    dus_cohort = dus_cohort,
    observation_period = observation_period
  )

  operations <- list(
    "addInObservation" = addInObservation,
    "addCohortCountAttr" = addCohortCountAttr,
    "addSex" = addSex,
    "addAge" = addAge,
    "addCdmName" = addCdmName,
    "addCohortIntersectDays" = function(x) addCohortIntersectDays(
      x,
      targetCohortId = 1,
      targetDate = "cohort_start_date",
      targetCohortTable = "cohort2"
    ),
    "addCohortIntersectDate" = function(x) addCohortIntersectDate(
      x,
      targetCohortId = 1,
      targetDate = "cohort_start_date",
      targetCohortTable = "cohort2"
    ),
    "addCohortIntersectCount" = function(x) addCohortIntersectCount(x, targetCohortTable = "cohort2"),
    "addCohortIntersectFlag" = function(x) addCohortIntersectFlag(x, targetCohortTable = "cohort2"),
    "addCohortName" = addCohortName,
    "addDateOfBirth" = addDateOfBirth,
    "addDemographics" = addDemographics,
    "addFutureObservation" = addFutureObservation,
    "addTableIntersectCount" = function(x) addTableIntersectCount(x, tableName = "drug_exposure")
  )

  baseline_class <- class(cdm$cohort1)
  baseline_class <- baseline_class[baseline_class != "GeneratedCohortSet"]
  # Apply each operation to cdm$cohort1 and check the class consistency
  for (op_name in names(operations)) {
    op <- operations[[op_name]]
    result <- op(cdm$cohort1)
    class(result) <- class(result)[class(result) != "GeneratedCohortSet"]
    expect_identical(class(result), baseline_class,
                     info = paste("Testing operation:", op_name))
  }

  result_with_sequence <- cdm$cohort1 |>
    addSex() |>
    addAge() |>
    addCategories(
      variable = "age",
      categories = list("age_group" = list(c(0, 120))),
      overlap = TRUE
    )
  class(result_with_sequence) <- class(result_with_sequence)[class(result_with_sequence) != "GeneratedCohortSet"]

  expect_identical(class(result_with_sequence), baseline_class,
                   info = "Testing sequence with addSex, addAge, and addCategories")

})