diff --git a/ccx_messaging/watchers/stats_watcher.py b/ccx_messaging/watchers/stats_watcher.py index 218343d..f28c1c4 100644 --- a/ccx_messaging/watchers/stats_watcher.py +++ b/ccx_messaging/watchers/stats_watcher.py @@ -16,8 +16,10 @@ import logging import os +import tarfile import time +from insights.core.archives import TarExtractor, ZipExtractor from insights_messaging.watchers import EngineWatcher from prometheus_client import Counter, Histogram, start_http_server, REGISTRY @@ -136,9 +138,29 @@ def on_filter(self): LOG.debug("Receiving 'on_filter' callback") self._filtered_total.inc() - def on_extract(self, ctx, broker, extraction): - """On extract event handler.""" + def on_extract( + self, ctx, broker, extraction: tarfile.TarFile | TarExtractor | ZipExtractor + ) -> None: + """On extract event handler for extractor using engines.""" LOG.debug("Receiving 'on_extract' callback") + + if isinstance(extraction, tarfile.TarFile): + self.on_extract_with_tarfile(extraction) + + else: + self.on_extract_with_extractor(extraction) + + def on_extract_with_tarfile(self, extraction: tarfile.TarFile) -> None: + """On extract event handler for engines using tarfile.""" + tarfile_contents = extraction.getnames() + + if "openshift_lightspeed.json" in tarfile_contents: + self._archive_metadata["type"] = "ols" + elif os.path.join("config", "infrastructure.json") in tarfile_contents: + self._archive_metadata["type"] = "hypershift" + + def on_extract_with_extractor(self, extraction: TarExtractor | ZipExtractor) -> None: + """On extract event handler for engines using extractor.""" # Set archive_type label based on found file if os.path.exists(os.path.join(extraction.tmp_dir, "openshift_lightspeed.json")): self._archive_metadata["type"] = "ols" diff --git a/test/watchers/stats_watcher_test.py b/test/watchers/stats_watcher_test.py index f8d1397..8870fa6 100644 --- a/test/watchers/stats_watcher_test.py +++ b/test/watchers/stats_watcher_test.py @@ -14,10 +14,12 @@ """Module containing unit tests for the `ConsumerWatcher` class.""" -from unittest.mock import patch +import time +from tarfile import TarFile +from unittest.mock import MagicMock, patch import pytest -import time +from insights.core.archives import TarExtractor from ccx_messaging.watchers.stats_watcher import ( StatsWatcher, @@ -340,3 +342,47 @@ def test_reset_archive_metadata(label_value): w._reset_archive_metadata() assert w._archive_metadata["type"] == "ocp" + + +ON_EXTRACT_WITH_EXTRACTOR_IDENTIFIERS = [ + ("ols", [True, False]), + ("hypershift", [False, True]), + ("ocp", [False, False]), +] + + +@patch("ccx_messaging.watchers.stats_watcher.start_http_server", lambda *args: None) +@pytest.mark.parametrize("type_, exists_side_effect", ON_EXTRACT_WITH_EXTRACTOR_IDENTIFIERS) +def test_on_extract_with_extractor(type_, exists_side_effect): + """Test the method on_extract.""" + # construct watcher object + w = StatsWatcher(prometheus_port=8009) + + extraction_mock = MagicMock(spec=TarExtractor) + extraction_mock.tmp_dir = "" + + with patch("os.path.exists") as exists_mock: + exists_mock.side_effect = exists_side_effect + + w.on_extract(None, None, extraction_mock) + assert w._archive_metadata["type"] == type_ + + +ON_EXTRACT_WITH_TARFILE_IDENTIFIERS = [ + ("ols", ["openshift_lightspeed.json"]), + ("hypershift", ["config/id", "config/infrastructure.json"]), + ("ocp", ["config/id"]), +] + + +@patch("ccx_messaging.watchers.stats_watcher.start_http_server", lambda *args: None) +@pytest.mark.parametrize("type_, file_content", ON_EXTRACT_WITH_TARFILE_IDENTIFIERS) +def test_on_extract_with_tarfile(type_, file_content): + """Test the method on_extract.""" + # construct watcher object + w = StatsWatcher(prometheus_port=8009) + + extraction_mock = MagicMock(spec=TarFile) + extraction_mock.getnames.return_value = file_content + w.on_extract(None, None, extraction_mock) + assert w._archive_metadata["type"] == type_