tests for dataset builder

3 files changed, 123 insertions, 3 deletions

diff --git a/tests/conftest.py b/tests/conftest.py
index 1c8d23c0..bb35dad1 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,12 +1,46 @@
+import json
+from importlib import resources
 from pathlib import Path
 
 import pytest
 
-import tests.data
+import tests.data.common
 from sec_certs.configuration import config
+from sec_certs.dataset import CPEDataset, CVEDataset
 
 
 @pytest.fixture(scope="module", autouse=True)
 def load_test_config():
-    pth = Path(tests.data.__path__[0]) / "settings_tests.yml"
-    config.load_from_yaml(pth)
+    with resources.path(tests.data.common, "settings_tests.yml") as path:
+        config.load_from_yaml(path)
+
+
+@pytest.fixture(scope="module")
+def cve_dataset_path() -> Path:
+    with resources.path(tests.data.common, "cve_dataset.json") as cve_dataset_path:
+        return cve_dataset_path
+
+
+@pytest.fixture(scope="module")
+def cpe_match_feed() -> dict:
+    with resources.open_text(tests.data.common, "cpe_match_feed.json") as handle:
+        data = json.load(handle)
+    return data
+
+
+@pytest.fixture(scope="module")
+def cve_dataset(cve_dataset_path: Path, cpe_match_feed: dict) -> CVEDataset:
+    cve_dataset = CVEDataset.from_json(cve_dataset_path)
+    cve_dataset.build_lookup_dict(cpe_match_feed)
+    return cve_dataset
+
+
+@pytest.fixture(scope="module")
+def cpe_dataset_path() -> Path:
+    with resources.path(tests.data.common, "cpe_dataset.json") as cpe_dataset_path:
+        return cpe_dataset_path
+
+
+@pytest.fixture(scope="module")
+def cpe_dataset(cpe_dataset_path: Path) -> CPEDataset:
+    return CPEDataset.from_json(cpe_dataset_path)
diff --git a/tests/data/settings_tests.yml b/tests/data/common/settings_tests.yml
index 645f11ab..645f11ab 100644
--- a/tests/data/settings_tests.yml
+++ b/tests/data/common/settings_tests.yml
diff --git a/tests/test_nvd_dataset_builder.py b/tests/test_nvd_dataset_builder.py
new file mode 100644
index 00000000..d887596e
--- /dev/null
+++ b/tests/test_nvd_dataset_builder.py
@@ -0,0 +1,86 @@
+from datetime import date, datetime, timedelta
+from importlib import resources
+from typing import Any
+
+import pytest
+
+import tests.data
+from sec_certs.configuration import config
+from sec_certs.dataset import CCDataset, CPEDataset, CVEDataset
+from sec_certs.utils.nvd_dataset_builder import (
+    CpeMatchNvdDatasetBuilder,
+    CpeNvdDatasetBuilder,
+    CveNvdDatasetBuilder,
+    NvdDatasetBuilder,
+)
+
+
+@pytest.fixture(autouse=True)
+def load_test_config():
+    with resources.path(tests.data, "settings_tests.yml") as path:
+        config.load_from_yaml(path)
+
+
+@pytest.mark.skip(reason="not yet implemented on the web")
+def test_cpe_download_from_seccerts():
+    cpe_dataset = CPEDataset.from_web()
+    assert len(cpe_dataset) > 100000
+    assert cpe_dataset.last_update_timestamp > datetime.now() - timedelta(days=28)
+
+
+@pytest.mark.skip(reason="not yet implemented on the web")
+def test_cve_download_from_seccerts():
+    cve_dataset = CVEDataset.from_web()
+    assert len(cve_dataset) > 100000
+    assert cve_dataset.last_update_timestamp > datetime.now() - timedelta(days=28)
+
+
+@pytest.mark.skip(reason="not yet implemented on the web")
+def test_cpe_match_download_from_seccerts():
+    cpe_match_dict = CCDataset._prepare_cpe_match_dict()
+    assert len(cpe_match_dict["match_strings"]) > 100000
+    assert datetime.fromisoformat(cpe_match_dict["timestamp"]) > datetime.now() - timedelta(days=28)
+
+
+@pytest.mark.parametrize(
+    "default_dataset, builder_class",
+    [
+        (CPEDataset(), CpeNvdDatasetBuilder),
+        (CVEDataset(), CveNvdDatasetBuilder),
+        ({"timestamp": datetime.fromtimestamp(0).isoformat(), "match_strings": {}}, CpeMatchNvdDatasetBuilder),
+    ],
+)
+def test_build_dataset(default_dataset: Any, builder_class: NvdDatasetBuilder):
+    def get_update_timestamp_from_dataset(dset) -> datetime:
+        if isinstance(dset, (CPEDataset, CVEDataset)):
+            return dset.last_update_timestamp
+        return datetime.fromisoformat(dset["timestamp"])
+
+    def get_dataset_len(dset) -> int:
+        if isinstance(dset, (CPEDataset, CVEDataset)):
+            return len(dset)
+        return len(dset["match_strings"])
+
+    config.preferred_source_nvd_datasets = "api"
+    with builder_class() as dataset_builder:
+        dataset = dataset_builder._init_new_dataset()
+        assert dataset == default_dataset
+        last_update = dataset_builder._get_last_update_from_previous_data(dataset)
+        dataset_builder._fill_in_mod_dates(False, last_update)
+        dataset_builder._build_arguments()
+
+        assert len(dataset_builder._requests_to_process) > 90
+        assert not dataset_builder._ok_responses
+        assert dataset_builder._attempts_left == dataset_builder.max_attempts
+
+        dataset_builder._requests_to_process = dataset_builder._requests_to_process[:1]
+        dataset_builder._request_parallel_and_handle_responses()
+        assert not dataset_builder._requests_to_process
+        assert len(dataset_builder._ok_responses) == 1
+
+        dataset = dataset_builder._process_responses(dataset_builder._ok_responses, dataset)
+
+        assert get_update_timestamp_from_dataset(dataset).date() == date.today()
+        assert (
+            get_dataset_len(dataset) > 200
+        )  # some items may be irrelevant, it's hard to tell how many, this is Bulgarian constant.