test_that("aug_bal_ate and aug_bal_att return accurate causal estimates without NAs or warnings", {
  set.seed(42)
  n <- 120
  d <- 4
  X <- matrix(rnorm(n * d), n, d)
  Z <- rbinom(n, 1, 0.5)
  Y <- 2 * Z + X %*% c(1, -0.5, 1, 0) + rnorm(n)

  fit_ate <- aug_bal_ate(Y, Z, X, type = "l2")
  expect_s3_class(fit_ate, "augbalweight")
  expect_true(is.numeric(fit_ate$estimate))
  expect_false(is.na(fit_ate$estimate))
  expect_false(is.nan(fit_ate$estimate))
  expect_true(fit_ate$se > 0)

  fit_att <- aug_bal_att(Y, Z, X, type = "l2")
  expect_s3_class(fit_att, "augbalweight")
  expect_true(is.numeric(fit_att$estimate))
  expect_false(is.na(fit_att$estimate))
  expect_false(is.nan(fit_att$estimate))
  expect_true(fit_att$se > 0)
})