From d7ec3624fdea1a5cbbf3e6933b965dcd9d8ca6b8 Mon Sep 17 00:00:00 2001
From: RayPlante <raymond.plante@nist.gov>
Date: Mon, 8 Jan 2024 15:47:59 -0500
Subject: [PATCH] integrate HeadBagRestorer, updated HeadBagCacheManager into
 HybridPDRDatasetRestorer

---
 .../distrib/cachemgr/pdr/HybridPDRDatasetRestorer.java   | 9 +++++++++
 .../cachemgr/pdr/HybridPDRDatasetRestorerTest.java       | 8 +++++---
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/main/java/gov/nist/oar/distrib/cachemgr/pdr/HybridPDRDatasetRestorer.java b/src/main/java/gov/nist/oar/distrib/cachemgr/pdr/HybridPDRDatasetRestorer.java
index 33d3111d..0cabcc34 100644
--- a/src/main/java/gov/nist/oar/distrib/cachemgr/pdr/HybridPDRDatasetRestorer.java
+++ b/src/main/java/gov/nist/oar/distrib/cachemgr/pdr/HybridPDRDatasetRestorer.java
@@ -172,6 +172,15 @@ public void restoreObject(String id, Reservation resv, String name, JSONObject m
         }
     }
 
+    @Override
+    protected InputStream openBag(String bagfilename) throws FileNotFoundException, StorageVolumeException {
+        try {
+            return ltstore.openFile(bagfilename);
+        } catch (FileNotFoundException ex) {
+            return restrictedLtstore.openFile(bagfilename);
+        }
+    }
+
     /**
      * consult the given head bag and return the name of the bag that contains the file indicated by
      * the given filepath.
diff --git a/src/test/java/gov/nist/oar/distrib/cachemgr/pdr/HybridPDRDatasetRestorerTest.java b/src/test/java/gov/nist/oar/distrib/cachemgr/pdr/HybridPDRDatasetRestorerTest.java
index cedac526..95782062 100644
--- a/src/test/java/gov/nist/oar/distrib/cachemgr/pdr/HybridPDRDatasetRestorerTest.java
+++ b/src/test/java/gov/nist/oar/distrib/cachemgr/pdr/HybridPDRDatasetRestorerTest.java
@@ -53,7 +53,9 @@ public class HybridPDRDatasetRestorerTest {
     ConfigurableCache cache = null;
 
     // Helper method to create HeadBagCacheManager. Adjust based on your actual implementation.
-    HeadBagCacheManager createHBCache(BagStorage ltstore) throws IOException, CacheManagementException {
+    HeadBagCacheManager createHBCache(BagStorage reststore, BagStorage pubstore)
+        throws IOException, CacheManagementException
+    {
         File tf = tempf.newFolder("headbags");
         File dbf = new File(tf, "inventory.sqlite");
         HeadBagDB.initializeSQLiteDB(dbf.getAbsolutePath());
@@ -66,7 +68,7 @@ HeadBagCacheManager createHBCache(BagStorage ltstore) throws IOException, CacheM
         cvd = new File(tf, "cv1");  cvd.mkdir();
         cache.addCacheVolume(new FilesystemCacheVolume(cvd, "cv1"), 2000000, null, true);
 
-        return new HeadBagCacheManager(cache, sidb, ltstore, new FileCopyRestorer(ltstore), "88434");
+        return new HeadBagCacheManager(cache, sidb, new HeadBagRestorer(reststore, pubstore), "88434");
     }
 
     ConfigurableCache createDataCache() throws CacheManagementException, IOException {
@@ -102,7 +104,7 @@ public void setUp() throws IOException, CacheManagementException {
         publicLtstore = new FilesystemLongTermStorage(ltsdir);
         restrictedLtstore = new FilesystemLongTermStorage(ltsdir + "/restricted");
 
-        hbcm = createHBCache(restrictedLtstore);
+        hbcm = createHBCache(restrictedLtstore, publicLtstore);
         cache = createDataCache();
         rstr = new HybridPDRDatasetRestorer(publicLtstore, restrictedLtstore, hbcm, 500);