From ba3e22a6b9d80db19abeb96aa08d48d26a3cb664 Mon Sep 17 00:00:00 2001
From: Katie Mummah <radioactivekate@gmail.com>
Date: Tue, 16 Jul 2024 17:14:36 -0500
Subject: [PATCH] bug fix and transport test

---
 CHANGELOG.rst       |  1 +
 src/source.cc       |  3 +--
 src/source_tests.cc | 39 +++++++++++++++++++++++++++++++++++----
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index bbb17d391..0a67ead68 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -7,6 +7,7 @@ Since last release
 
 **Added:**
 * Added package parameter to storage (#603, #612)
+* Added package parameter to source (#613)
 
 **Changed:**
 
diff --git a/src/source.cc b/src/source.cc
index 48a80274d..9844318a7 100644
--- a/src/source.cc
+++ b/src/source.cc
@@ -109,7 +109,7 @@ std::set<cyclus::BidPortfolio<cyclus::Material>::Ptr> Source::GetMatlBids(
     std::vector<double> bids;
     bids.assign(n_full_bids, bid_qty);
 
-    double remaining_qty = std::fmod(qty, bid_qty);
+    double remaining_qty = qty - (n_full_bids * bid_qty);
     if ((remaining_qty > cyclus::eps()) && (remaining_qty >= context()->GetPackage(package)->fill_min())) {
       bids.push_back(remaining_qty);
     }
@@ -165,7 +165,6 @@ void Source::GetMatlTrades(
         // If not all material is packaged successfully, return the excess
         // amount to the inventory
         inventory.Push(m);
-        //inventory_size += m->quantity();
       }
 
       Material::Ptr response;
diff --git a/src/source_tests.cc b/src/source_tests.cc
index ed518cdb1..42c6fe551 100644
--- a/src/source_tests.cc
+++ b/src/source_tests.cc
@@ -172,8 +172,6 @@ TEST_F(SourceTest, Package) {
   
   sim.context()->AddPackage(package_name, 3, 4, "first");
   package = sim.context()->GetPackage(package_name);
-  // sim.context()->AddTransportUnit(tu_name, 1, 2, "hybrid");
-  // tu = sim.context()->GetTransportUnit(tu_name);
   
   sim.AddSink("commod").Finalize();
 
@@ -186,8 +184,41 @@ TEST_F(SourceTest, Package) {
   conds.push_back(Cond("PackageName", "==", package->name()));
   QueryResult qr_res = sim.db().Query("Resources", &conds);
 
-  EXPECT_EQ(qr_res.rows.size(), 3);
-    
+  EXPECT_EQ(qr_res.rows.size(), 3); 
+}
+
+TEST_F(SourceTest, TransportUnit) {
+  using cyclus::QueryResult;
+  using cyclus::Cond;
+
+  std::string config =
+    "<outcommod>commod</outcommod>"
+    "<package>testpackage</package>"
+    "<transport_unit>testtu</transport_unit>"
+    "<throughput>10</throughput>";
+
+  int simdur = 2;
+  cyclus::MockSim sim(cyclus::AgentSpec (":cycamore:Source"), config, simdur);
+  
+  sim.context()->AddPackage(package_name, 3, 4, "equal");
+  package = sim.context()->GetPackage(package_name);
+  sim.context()->AddTransportUnit(tu_name, 2, 2);
+  tu = sim.context()->GetTransportUnit(tu_name);
+  
+  sim.AddSink("commod").Finalize();
+
+  EXPECT_NO_THROW(sim.Run());
+
+  QueryResult qr_tr = sim.db().Query("Transactions", NULL);
+  EXPECT_EQ(qr_tr.rows.size(), 4);
+  
+  std::vector<Cond> conds;
+  conds.push_back(Cond("PackageName", "==", package->name()));
+  QueryResult qr_res = sim.db().Query("Resources", &conds);
+
+  EXPECT_EQ(qr_res.rows.size(), 4);
+
+  QueryResult qr_allres = sim.db().Query("Resources", NULL);
 }
 
 boost::shared_ptr< cyclus::ExchangeContext<cyclus::Material> >