fix: test_output_json2 failures in test_output_engine.py (#4903)

Author: Saksham-Sirohi

cve_bin_tool/output_engine/json_output.py

@@ -8,6 +8,7 @@
 from typing import IO
 
 from cve_bin_tool.cvedb import CVEDB
+from cve_bin_tool.log import LOGGER
 from cve_bin_tool.util import CVEData, ProductInfo, VersionInfo
 from cve_bin_tool.version import VERSION
 
@@ -45,20 +46,50 @@ def db_entries_count():
     """
     Retrieves the count of CVE entries from the database grouped by data source.
 
+    The function fetches all data sources from the database and ensures that
+    schema-required sources (NVD, OSV, GAD, REDHAT) are always included in the result,
+    even if they don't exist in the database yet, while also including all other
+    sources with counts initialized to 0.
+
     Returns:
         dict: A dictionary containing the count of CVE entries for each data source.
     """
     instance = CVEDB()
     cursor = instance.db_open_and_get_cursor()
-    cve_entries_check = "SELECT data_source, COUNT(*) as number FROM cve_severity GROUP BY data_source ORDER BY number DESC"
-    cursor.execute(cve_entries_check)
     data_entries = {}
-    rows = cursor.fetchall()
-    for row in rows:
-        source = row[0]
-        entries = row[1]
-        data_entries[source] = entries
-    instance.db_close()
+
+    # Always include the schema-required sources
+    required_sources = ["NVD", "OSV", "GAD", "REDHAT"]
+    for source in required_sources:
+        data_entries[source] = 0
+
+    try:
+        # Get all possible sources from the database
+        sources_query = "SELECT DISTINCT data_source FROM cve_severity"
+        cursor.execute(sources_query)
+        all_sources = [row[0] for row in cursor.fetchall()]
+
+        # Initialize counts to 0 for all sources not in required_sources
+        for source in all_sources:
+            if source not in data_entries:
+                data_entries[source] = 0
+
+        # Get actual counts
+        cve_entries_check = (
+            "SELECT data_source, COUNT(*) as number FROM cve_severity "
+            "GROUP BY data_source ORDER BY number DESC"
+        )
+        cursor.execute(cve_entries_check)
+        rows = cursor.fetchall()
+        for row in rows:
+            source = row[0]
+            entries = row[1]
+            data_entries[source] = entries
+
+    except Exception as e:
+        LOGGER.error(f"Error in db_entries_count or cve_severity in db: {e}")
+    finally:
+        instance.db_close()
     return data_entries