From de3a58d1054b3fdd1137fadae3ed8029732b91a3 Mon Sep 17 00:00:00 2001
From: Jared Knowles <jared@civilytics.com>
Date: Mon, 3 Aug 2026 10:19:32 -0400
Subject: [PATCH] fix: attach govids_found/govids_missing to cog_balances()
 provenance

Mirrors R/spending.R:465-466 -- .check_govids_in_scope()'s return was
previously captured only for its message side effect. Also drops a
redundant duplicate assertion in the flow-code guard test.
---
 R/balances.R                   |  4 +++-
 tests/testthat/test-balances.R | 13 ++++++++++++-
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/R/balances.R b/R/balances.R
index fd8b4ee..bfd6e77 100644
--- a/R/balances.R
+++ b/R/balances.R
@@ -55,7 +55,7 @@ cog_balances <- function(govid, years, category = NULL,
 
   con <- .ensure_session()
   .require_balance_support(con)
-  .check_govids_in_scope(govid)
+  scope <- .check_govids_in_scope(govid)
 
   basis_note <- paste0(
     "`basis` has no effect on holdings: harmonization_map carries no ",
@@ -77,6 +77,8 @@ cog_balances <- function(govid, years, category = NULL,
     expenditure_concept = NA_character_,
     revenue_concept = NA_character_
   )
+  prov$scope$govids_found   <- scope$found
+  prov$scope$govids_missing <- scope$missing
 
   attr(result, "provenance") <- prov
   result
diff --git a/tests/testthat/test-balances.R b/tests/testthat/test-balances.R
index a5f6efa..2313194 100644
--- a/tests/testthat/test-balances.R
+++ b/tests/testthat/test-balances.R
@@ -144,7 +144,6 @@ test_that("no flow code can reach cog_balances", {
 
     expect_true(length(got) > 0L)
     expect_true(all(got %in% balance_codes))
-    expect_true(length(setdiff(got, balance_codes)) == 0L)
   })
 })
 
@@ -166,3 +165,15 @@ test_that("every balance_subtype maps to exactly one category", {
                         function(x) length(unique(x)))
   expect_true(all(per_subtype == 1L))
 })
+
+test_that("cog_balances records found + missing govids in provenance", {
+  skip_if_no_corpus()
+  with_fixture_corpus({
+    suppressMessages(
+      r <- cog_balances(c("550000227544", "XXXINVALID"), 2019)
+    )
+    prov <- attr(r, "provenance")
+    expect_equal(sort(prov$scope$govids_found), "550000227544")
+    expect_equal(sort(prov$scope$govids_missing), "XXXINVALID")
+  })
+})