add aux dataset tests

5 files changed, 313 insertions, 92 deletions

diff --git a/docs/user_guide.md b/docs/user_guide.md
index 9d330c6c..8e6e2f4b 100644
--- a/docs/user_guide.md
+++ b/docs/user_guide.md
@@ -16,11 +16,11 @@ Our tool can seamlessly download the required NVD datasets when needed. We suppo
 The following two keys control the behaviour:
 
 ```yaml
-preferred_source_nvd_datasets: "api" # set to "sec-certs" to fetch them from sec-certs.org
+preferred_source_aux_datasets: "api" # set to "sec-certs" to fetch them from sec-certs.org
 nvd_api_key: null # or the actual key value
 ```
 
-If you aim to fetch the sources from NVD, we advise you to get an [NVD API key](https://nvd.nist.gov/developers/request-an-api-key) and set the `nvd_api_key` setting accordingly. The download from NVD will work even without API key, it will just be slow. No API key is needed when `preferred_source_nvd_datasets: "sec-certs"`
+If you aim to fetch the sources from NVD, we advise you to get an [NVD API key](https://nvd.nist.gov/developers/request-an-api-key) and set the `nvd_api_key` setting accordingly. The download from NVD will work even without API key, it will just be slow. No API key is needed when `preferred_source_aux_datasets: "sec-certs"`
 
 
 ## Inferring inter-certificate reference context
diff --git a/src/sec_certs/configuration.py b/src/sec_certs/configuration.py
index b12dfb07..1bf5df1a 100644
--- a/src/sec_certs/configuration.py
+++ b/src/sec_certs/configuration.py
@@ -138,7 +138,7 @@ class Configuration(BaseSettings):
         description="If true, progress bars will be printed to stdout during computation.",
     )
     nvd_api_key: Optional[str] = Field(None, description="NVD API key for access to CVEs and CPEs.")  # noqa: UP007
-    preferred_source_nvd_datasets: Literal["sec-certs", "api"] = Field(
+    preferred_source_aux_datasets: Literal["sec-certs", "api"] = Field(
         "sec-certs",
         description="If set to `sec-certs`, will fetch CPE and CVE datasets from sec-certs.org."
         + " If set to `api`, will fetch these resources from NVD API. It is advised to set an"
diff --git a/src/sec_certs/dataset/auxiliary_dataset_handling.py b/src/sec_certs/dataset/auxiliary_dataset_handling.py
index 829ecf07..0704df2a 100644
--- a/src/sec_certs/dataset/auxiliary_dataset_handling.py
+++ b/src/sec_certs/dataset/auxiliary_dataset_handling.py
@@ -64,22 +64,21 @@ class CPEDatasetHandler(AuxiliaryDatasetHandler):
 
     @staged(logger, "Processing CPE dataset")
     def _process_dataset_body(self, download_fresh: bool = False) -> None:
-        if self.dset_path.exists():
+        if not download_fresh and self.dset_path.exists():
             logger.info("Preparing CPEDataset from json.")
-            self.dset = CPEDataset.from_json(self.dset_path)
+            self.load_dataset()
+            return
+
+        if config.preferred_source_aux_datasets == "api":
+            logger.info("Fetching new CPE records from NVD API")
+            with CpeNvdDatasetBuilder(api_key=config.nvd_api_key) as builder:
+                self.dset = builder.build_dataset()
         else:
-            self.dset = CPEDataset(json_path=self.dset_path)
-            download_fresh = True
+            logger.info("Preparing CPEDataset from sec-certs.org.")
+            self.dset = CPEDataset.from_web(self.dset_path)
 
-        if download_fresh:
-            if config.preferred_source_nvd_datasets == "api":
-                logger.info("Fetching new CPE records from NVD API")
-                with CpeNvdDatasetBuilder(api_key=config.nvd_api_key) as builder:
-                    self.dset = builder.build_dataset(self.dset)
-            else:
-                logger.info("Preparing CPEDataset from sec-certs.org.")
-                self.dset = CPEDataset.from_web(self.dset_path)
-            self.dset.to_json()
+        self.dset.to_json()
+        self.dset.json_path = self.dset_path
 
     def load_dataset(self) -> None:
         self.dset = CPEDataset.from_json(self.dset_path)
@@ -92,22 +91,21 @@ class CVEDatasetHandler(AuxiliaryDatasetHandler):
 
     @staged(logger, "Processing CVE dataset")
     def _process_dataset_body(self, download_fresh: bool = False) -> None:
-        if self.dset_path.exists():
+        if not download_fresh and self.dset_path.exists():
             logger.info("Preparing CVEDataset from json.")
-            self.dset = CVEDataset.from_json(self.dset_path)
+            self.load_dataset()
+            return
+
+        if config.preferred_source_aux_datasets == "api":
+            logger.info("Fetching new CVE records from NVD API.")
+            with CveNvdDatasetBuilder(api_key=config.nvd_api_key) as builder:
+                self.dset = builder.build_dataset()
         else:
-            self.dset = CVEDataset(json_path=self.dset_path)
-            download_fresh = True
+            logger.info("Preparing CVEDataset from sec-certs.org.")
+            self.dset = CVEDataset.from_web(self.dset_path)
 
-        if download_fresh:
-            if config.preferred_source_nvd_datasets == "api":
-                logger.info("Fetching new CVE records from NVD API.")
-                with CveNvdDatasetBuilder(api_key=config.nvd_api_key) as builder:
-                    self.dset = builder.build_dataset(self.dset)
-            else:
-                logger.info("Preparing CVEDataset from sec-certs.org")
-                self.dset = CVEDataset.from_web(self.dset_path)
-            self.dset.to_json()
+        self.dset.to_json()
+        self.dset.json_path = self.dset_path
 
     def load_dataset(self):
         self.dset = CVEDataset.from_json(self.dset_path)
@@ -120,39 +118,34 @@ class CPEMatchDictHandler(AuxiliaryDatasetHandler):
 
     @staged(logger, "Processing CPE Match dictionary")
     def _process_dataset_body(self, download_fresh: bool = False) -> None:
-        if self.dset_path.exists():
+        if not download_fresh and self.dset_path.exists():
             logger.info("Preparing CPE Match feed from json.")
-            with self.dset_path.open("r") as handle:
-                self.dset = json.load(handle)
+            self.load_dataset()
+            return
+
+        if config.preferred_source_aux_datasets == "api":
+            logger.info("Fetchnig CPE Match feed from NVD APi.")
+            with CpeMatchNvdDatasetBuilder(api_key=config.nvd_api_key) as builder:
+                self.dset = builder.build_dataset()
         else:
-            self.dset = CpeMatchNvdDatasetBuilder._init_new_dataset()
-            download_fresh = True
+            logger.info("Preparing CPE Match feed from sec-certs.org.")
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                dset_path = Path(tmp_dir) / "cpe_match_feed.json.gz"
+                if (
+                    not helpers.download_file(
+                        config.cpe_match_latest_snapshot,
+                        dset_path,
+                        progress_bar_desc="Downloading CPE Match feed from web",
+                    )
+                    == constants.RESPONSE_OK
+                ):
+                    raise RuntimeError(f"Could not download CPE Match feed from {config.cpe_match_latest_snapshot}.")
+                with gzip.open(str(dset_path)) as handle:
+                    json_str = handle.read().decode("utf-8")
+                    self.dset = json.loads(json_str)
 
-        if download_fresh:
-            if config.preferred_source_nvd_datasets == "api":
-                logger.info("Fetchnig CPE Match feed from NVD APi.")
-                with CpeMatchNvdDatasetBuilder(api_key=config.nvd_api_key) as builder:
-                    self.dset = builder.build_dataset(self.dset)
-            else:
-                logger.info("Preparing CPE Match feed from sec-certs.org.")
-                with tempfile.TemporaryDirectory() as tmp_dir:
-                    dset_path = Path(tmp_dir) / "cpe_match_feed.json.gz"
-                    if (
-                        not helpers.download_file(
-                            config.cpe_match_latest_snapshot,
-                            dset_path,
-                            progress_bar_desc="Downloading CPE Match feed from web",
-                        )
-                        == constants.RESPONSE_OK
-                    ):
-                        raise RuntimeError(
-                            f"Could not download CPE Match feed from {config.cpe_match_latest_snapshot}."
-                        )
-                    with gzip.open(str(dset_path)) as handle:
-                        json_str = handle.read().decode("utf-8")
-                        self.dset = json.loads(json_str)
-            with self.dset_path.open("w") as handle:
-                json.dump(self.dset, handle, indent=4)
+        with self.dset_path.open("w") as handle:
+            json.dump(self.dset, handle, indent=4)
 
     def load_dataset(self):
         with self.dset_path.open("r") as handle:
@@ -166,11 +159,14 @@ class FIPSAlgorithmDatasetHandler(AuxiliaryDatasetHandler):
 
     @staged(logger, "Processing FIPS Algorithms")
     def _process_dataset_body(self, download_fresh: bool = False) -> None:
-        if not self.dset_path.exists() or download_fresh:
-            self.dset = FIPSAlgorithmDataset.from_web(self.dset_path)
-            self.dset.to_json()
-        else:
-            self.dset = FIPSAlgorithmDataset.from_json(self.dset_path)
+        if not download_fresh and self.dset_path.exists():
+            logger.info("Preparing FIPSAlgorithmDataset from json.")
+            self.load_dataset()
+            return
+
+        self.dset = FIPSAlgorithmDataset.from_web(self.dset_path)
+        self.dset.to_json()
+        self.dset.json_path = self.dset_path
 
     def load_dataset(self):
         self.dset = FIPSAlgorithmDataset.from_json(self.dset_path)
@@ -192,11 +188,14 @@ class CCSchemeDatasetHandler(AuxiliaryDatasetHandler):
 
     @staged(logger, "Processing CC Schemes")
     def _process_dataset_body(self, download_fresh: bool = False) -> None:
-        if not self.dset_path.exists() or download_fresh:
-            self.dset = CCSchemeDataset.from_web(self.dset_path, self.only_schemes)
-            self.dset.to_json()
-        else:
-            self.dset = CCSchemeDataset.from_json(self.dset_path)
+        if not download_fresh and self.dset_path.exists():
+            logger.info("Preparing CCSchemeDataset from json.")
+            self.load_dataset()
+            return
+
+        self.dset = CCSchemeDataset.from_web(self.dset_path, self.only_schemes)
+        self.dset.to_json()
+        self.dset.json_path = self.dset_path
 
     def load_dataset(self):
         self.dset = CCSchemeDataset.from_json(self.dset_path)
@@ -227,20 +226,25 @@ class CCMaintenanceUpdateDatasetHandler(AuxiliaryDatasetHandler):
     def _process_dataset_body(self, download_fresh: bool = False):
         from sec_certs.dataset.cc import CCDatasetMaintenanceUpdates
 
-        if not self.dset_path.exists() or download_fresh:
-            updates = list(
-                itertools.chain.from_iterable(
-                    CCMaintenanceUpdate.get_updates_from_cc_cert(x) for x in self.certs_with_updates
-                )
-            )
-            self.dset = CCDatasetMaintenanceUpdates(
-                {x.dgst: x for x in updates}, root_dir=self.dset_path.parent, name="maintenance_updates"
+        if not download_fresh and self.dset_path.exists():
+            logger.info("Preparing CCDatasetMaintenanceUpdates from json.")
+            self.load_dataset()
+            return
+
+        updates = list(
+            itertools.chain.from_iterable(
+                CCMaintenanceUpdate.get_updates_from_cc_cert(x) for x in self.certs_with_updates
             )
-            self.dset.download_all_artifacts()
-            self.dset.convert_all_pdfs()
-            self.dset.extract_data()
-        else:
-            self.dset = CCDatasetMaintenanceUpdates.from_json(self.dset_path)
+        )
+        self.dset = CCDatasetMaintenanceUpdates(
+            {x.dgst: x for x in updates},
+            root_dir=self.dset_path.parent,
+            name="maintenance_updates",
+        )
+        self.dset.download_all_artifacts()
+        self.dset.convert_all_pdfs()
+        self.dset.extract_data()
+        self.dset.to_json()
 
 
 class ProtectionProfileDatasetHandler(AuxiliaryDatasetHandler):
@@ -262,12 +266,14 @@ class ProtectionProfileDatasetHandler(AuxiliaryDatasetHandler):
     def _process_dataset_body(self, download_fresh: bool = False):
         from sec_certs.dataset.protection_profile import ProtectionProfileDataset
 
-        if not self.dset_path.exists() or download_fresh:
-            self.dset_path.parent.mkdir(exist_ok=True, parents=True)
-            self.dset = ProtectionProfileDataset(root_dir=self.dset_path.parent)
-            self.dset.get_certs_from_web()
-            self.dset.download_all_artifacts()
-            self.dset.convert_all_pdfs()
-            self.dset.analyze_certificates()
-        else:
-            self.dset = ProtectionProfileDataset.from_json(self.dset_path)
+        if not download_fresh and self.dset_path.exists():
+            logger.info("Preparing ProtectionProfileDataset from json.")
+            self.load_dataset()
+            return
+
+        self.dset_path.parent.mkdir(exist_ok=True, parents=True)
+        self.dset = ProtectionProfileDataset(root_dir=self.dset_path.parent)
+        self.dset.get_certs_from_web()
+        self.dset.download_all_artifacts()
+        self.dset.convert_all_pdfs()
+        self.dset.analyze_certificates()
diff --git a/tests/cc/test_cc_aux_datasets.py b/tests/cc/test_cc_aux_datasets.py
new file mode 100644
index 00000000..8a7a32a9
--- /dev/null
+++ b/tests/cc/test_cc_aux_datasets.py
@@ -0,0 +1,215 @@
+from unittest.mock import mock_open
+
+import pytest
+
+from sec_certs.configuration import config
+from sec_certs.dataset import (
+    CCDatasetMaintenanceUpdates,
+    CCSchemeDataset,
+    CPEDataset,
+    CVEDataset,
+    FIPSAlgorithmDataset,
+    ProtectionProfileDataset,
+)
+from sec_certs.dataset.auxiliary_dataset_handling import (
+    CCMaintenanceUpdateDatasetHandler,
+    CCSchemeDatasetHandler,
+    CPEDatasetHandler,
+    CPEMatchDictHandler,
+    CVEDatasetHandler,
+    FIPSAlgorithmDatasetHandler,
+    ProtectionProfileDatasetHandler,
+)
+
+
+@pytest.fixture
+def temp_dir(tmp_path):
+    return tmp_path
+
+
+@pytest.fixture
+def mock_dset():
+    return {"key": "value"}
+
+
+def test_cpe_dataset_handler_set_local_paths(temp_dir):
+    handler = CPEDatasetHandler(temp_dir)
+    new_path = temp_dir / "new_path"
+    handler.set_local_paths(new_path)
+    assert handler.aux_datasets_dir == new_path
+
+
+@pytest.mark.parametrize("preferred_source_aux_datasets", ["sec-certs", "api"])
+def test_cpe_dataset_handler_process_dataset(preferred_source_aux_datasets, temp_dir, monkeypatch):
+    config.preferred_source_aux_datasets = preferred_source_aux_datasets
+    handler = CPEDatasetHandler(temp_dir)
+    mock_dset = CPEDataset()
+
+    def mock_get_dset(path):
+        return mock_dset
+
+    if preferred_source_aux_datasets == "sec-certs":
+        monkeypatch.setattr("sec_certs.dataset.cpe.CPEDataset.from_web", mock_get_dset)
+    else:
+        monkeypatch.setattr("sec_certs.utils.nvd_dataset_builder.CpeNvdDatasetBuilder.build_dataset", mock_get_dset)
+
+    monkeypatch.setattr("sec_certs.dataset.cpe.CPEDataset.to_json", lambda x: None)
+    handler.process_dataset(download_fresh=True)
+
+    assert handler.dset == mock_dset
+    assert handler.dset_path == temp_dir / "cpe_dataset.json"
+
+
+def test_cve_dataset_handler_set_local_paths(temp_dir):
+    handler = CVEDatasetHandler(temp_dir)
+    new_path = temp_dir / "new_path"
+    handler.set_local_paths(new_path)
+    assert handler.aux_datasets_dir == new_path
+
+
+@pytest.mark.parametrize("preferred_source_aux_datasets", ["sec-certs", "api"])
+def test_cve_dataset_handler_process_dataset(preferred_source_aux_datasets, temp_dir, monkeypatch):
+    config.preferred_source_aux_datasets = preferred_source_aux_datasets
+    handler = CVEDatasetHandler(temp_dir)
+    mock_dset = CVEDataset()
+
+    def mock_get_dset(path):
+        return mock_dset
+
+    if preferred_source_aux_datasets == "sec-certs":
+        monkeypatch.setattr("sec_certs.dataset.cve.CVEDataset.from_web", mock_get_dset)
+    else:
+        monkeypatch.setattr("sec_certs.utils.nvd_dataset_builder.CveNvdDatasetBuilder.build_dataset", mock_get_dset)
+    monkeypatch.setattr("sec_certs.dataset.cve.CVEDataset.to_json", lambda x: None)
+    handler.process_dataset(download_fresh=True)
+
+    assert handler.dset == mock_dset
+    assert handler.dset_path == temp_dir / "cve_dataset.json"
+
+
+def test_cpe_match_dict_handler_set_local_paths(temp_dir):
+    handler = CPEMatchDictHandler(temp_dir)
+    new_path = temp_dir / "new_path"
+    handler.set_local_paths(new_path)
+    assert handler.aux_datasets_dir == new_path
+
+
+@pytest.mark.parametrize("preferred_source_aux_datasets", ["sec-certs", "api"])
+def test_cpe_match_dict_handler_process_dataset(preferred_source_aux_datasets, temp_dir, monkeypatch):
+    config.preferred_source_aux_datasets = preferred_source_aux_datasets
+    handler = CPEMatchDictHandler(temp_dir)
+    mock_dset = {"key": "value"}
+    mock_dset_str_single_quotes = '{"key": "value"}'
+
+    def mock_get_dset(path):
+        return mock_dset
+
+    def mock_download_file(url, path, progress_bar_desc):
+        return 200
+
+    if preferred_source_aux_datasets == "api":
+        monkeypatch.setattr(
+            "sec_certs.utils.nvd_dataset_builder.CpeMatchNvdDatasetBuilder.build_dataset", mock_get_dset
+        )
+    else:
+        monkeypatch.setattr("sec_certs.utils.helpers.download_file", mock_download_file)
+        monkeypatch.setattr("gzip.open", mock_open(read_data=(mock_dset_str_single_quotes.encode())))
+
+    handler.process_dataset(download_fresh=True)
+
+    assert handler.dset == mock_dset
+
+
+def test_fips_algorithm_dataset_handler_set_local_paths(temp_dir):
+    handler = FIPSAlgorithmDatasetHandler(temp_dir)
+    new_path = temp_dir / "new_path"
+    handler.set_local_paths(new_path)
+    assert handler.aux_datasets_dir == new_path
+
+
+def test_fips_algorithm_dataset_handler_process_dataset(temp_dir, monkeypatch):
+    handler = FIPSAlgorithmDatasetHandler(temp_dir)
+    mock_dset = FIPSAlgorithmDataset()
+
+    def mock_from_web(path):
+        return mock_dset
+
+    monkeypatch.setattr("sec_certs.dataset.fips_algorithm.FIPSAlgorithmDataset.from_web", mock_from_web)
+    monkeypatch.setattr("sec_certs.dataset.fips_algorithm.FIPSAlgorithmDataset.to_json", lambda x: None)
+    handler.process_dataset(download_fresh=True)
+    assert handler.dset == mock_dset
+    assert handler.dset_path == temp_dir / "algorithms.json"
+    assert handler.dset.json_path == handler.dset_path
+
+
+def test_cc_scheme_dataset_handler_set_local_paths(temp_dir):
+    handler = CCSchemeDatasetHandler(temp_dir)
+    new_path = temp_dir / "new_path"
+    handler.set_local_paths(new_path)
+    assert handler.aux_datasets_dir == new_path
+
+
+def test_cc_scheme_dataset_handler_process_dataset(temp_dir, monkeypatch):
+    handler = CCSchemeDatasetHandler(temp_dir)
+    mock_dset = CCSchemeDataset(schemes={})
+
+    def mock_from_web(path, only_schemes):
+        return mock_dset
+
+    monkeypatch.setattr("sec_certs.dataset.cc_scheme.CCSchemeDataset.from_web", mock_from_web)
+    monkeypatch.setattr("sec_certs.dataset.cc_scheme.CCSchemeDataset.to_json", lambda x: None)
+    handler.process_dataset(download_fresh=True)
+    assert handler.dset == mock_dset
+    assert handler.dset_path == temp_dir / "cc_scheme.json"
+    assert handler.dset.json_path == handler.dset_path
+
+
+def test_cc_maintenance_update_dataset_handler_set_local_paths(temp_dir):
+    handler = CCMaintenanceUpdateDatasetHandler(temp_dir)
+    new_path = temp_dir / "new_path"
+    handler.set_local_paths(new_path)
+    assert handler.aux_datasets_dir == new_path
+
+
+def test_cc_maintenance_update_dataset_handler_process_dataset(temp_dir, monkeypatch):
+    handler = CCMaintenanceUpdateDatasetHandler(temp_dir)
+    mock_dset = CCDatasetMaintenanceUpdates(root_dir=handler.dset_path.parent, name="maintenance_updates")
+
+    monkeypatch.setattr(
+        "sec_certs.sample.cc_maintenance_update.CCMaintenanceUpdate.get_updates_from_cc_cert",
+        lambda x: [],
+    )
+    monkeypatch.setattr("sec_certs.dataset.dataset.Dataset.download_all_artifacts", lambda x: None)
+    monkeypatch.setattr("sec_certs.dataset.dataset.Dataset.convert_all_pdfs", lambda x: None)
+    monkeypatch.setattr("sec_certs.dataset.cc.CCDataset.extract_data", lambda x: None)
+    monkeypatch.setattr("sec_certs.dataset.dataset.Dataset.to_json", lambda x: None)
+    handler.process_dataset(download_fresh=True)
+    assert handler.dset == mock_dset
+
+
+def test_protection_profile_dataset_handler_set_local_paths(temp_dir):
+    handler = ProtectionProfileDatasetHandler(temp_dir)
+    new_path = temp_dir / "new_path"
+    handler.set_local_paths(new_path)
+    assert handler.aux_datasets_dir == new_path
+
+
+def test_protection_profile_dataset_handler_process_dataset(temp_dir, monkeypatch):
+    handler = ProtectionProfileDatasetHandler(temp_dir)
+    mock_dset = ProtectionProfileDataset()
+
+    monkeypatch.setattr(
+        "sec_certs.dataset.protection_profile.ProtectionProfileDataset.get_certs_from_web", lambda x: None
+    )
+    monkeypatch.setattr(
+        "sec_certs.dataset.protection_profile.ProtectionProfileDataset.download_all_artifacts", lambda x: None
+    )
+    monkeypatch.setattr(
+        "sec_certs.dataset.protection_profile.ProtectionProfileDataset.convert_all_pdfs", lambda x: None
+    )
+    monkeypatch.setattr(
+        "sec_certs.dataset.protection_profile.ProtectionProfileDataset.analyze_certificates", lambda x: None
+    )
+    monkeypatch.setattr("sec_certs.dataset.protection_profile.ProtectionProfileDataset.to_json", lambda x: None)
+    handler.process_dataset(download_fresh=True)
+    assert handler.dset == mock_dset
diff --git a/tests/test_nvd_dataset_builder.py b/tests/test_nvd_dataset_builder.py
index fda94d1f..44ea7c44 100644
--- a/tests/test_nvd_dataset_builder.py
+++ b/tests/test_nvd_dataset_builder.py
@@ -60,7 +60,7 @@ def test_build_dataset(default_dataset: Any, builder_class: type[NvdDatasetBuild
             return len(dset)
         return len(dset["match_strings"])
 
-    config.preferred_source_nvd_datasets = "api"
+    config.preferred_source_aux_datasets = "api"
     with builder_class(api_key=config.nvd_api_key) as dataset_builder:
         dataset = dataset_builder._init_new_dataset()
         assert dataset == default_dataset