added utests

Author: alexandraBara

nodescraper/cli/helper.py

@@ -313,6 +313,33 @@ def extract_analyzer_args_from_model(
         return None
 
 
+def extract_collector_args_from_model(
+    plugin_cls: type, data_model: BaseModel, logger: logging.Logger
+) -> Optional[BaseModel]:
+    """Extract collector args from a plugin and a data model.
+
+    Args:
+        plugin_cls (type): The plugin class from registry.
+        data_model (BaseModel): System data model.
+        logger (logging.Logger): logger.
+
+    Returns:
+        Optional[BaseModel]: Instance of collector args model or None if unavailable.
+    """
+    if not hasattr(plugin_cls, "COLLECTOR_ARGS") or not plugin_cls.COLLECTOR_ARGS:
+        logger.warning(
+            "Plugin: %s does not support reference config creation. No collector args defined.",
+            getattr(plugin_cls, "__name__", str(plugin_cls)),
+        )
+        return None
+
+    try:
+        return plugin_cls.COLLECTOR_ARGS.build_from_model(data_model)
+    except NotImplementedError as e:
+        logger.info("%s: %s", plugin_cls.__name__, str(e))
+        return None
+
+
 def generate_reference_config(
     results: list[PluginResult], plugin_reg: PluginRegistry, logger: logging.Logger
 ) -> PluginConfig:
@@ -344,11 +371,16 @@ def generate_reference_config(
 
         plugin = plugin_reg.plugins.get(obj.source)
 
-        args = extract_analyzer_args_from_model(plugin, data_model, logger)
-        if not args:
-            continue
-        plugins[obj.source] = {"analysis_args": {}}
-        plugins[obj.source]["analysis_args"] = args.model_dump(exclude_none=True)
+        if obj.source not in plugins:
+            plugins[obj.source] = {}
+
+        a_args = extract_analyzer_args_from_model(plugin, data_model, logger)
+        if a_args:
+            plugins[obj.source]["analysis_args"] = a_args.model_dump(exclude_none=True)
+
+        c_args = extract_collector_args_from_model(plugin, data_model, logger)
+        if c_args:
+            plugins[obj.source]["collector_args"] = c_args.model_dump(exclude_none=True)
     plugin_config.plugins = plugins
 
     return plugin_config

nodescraper/models/analyzerargs.py

@@ -31,7 +31,7 @@ class AnalyzerArgs(BaseModel):
 
     @classmethod
     def build_from_model(cls, datamodel):
-        """Build analyzer args instance from data model object
+        """Build collector args instance from data model object
 
         Args:
             datamodel (TDataModel): data model to use for creating analyzer args
@@ -40,5 +40,5 @@ def build_from_model(cls, datamodel):
             NotImplementedError: Not implemented error
         """
         raise NotImplementedError(
-            "Setting analyzer args from datamodel is not implemented for class: %s", cls.__name__
+            "Setting collector args from datamodel is not implemented for class: %s", cls.__name__
         )

nodescraper/models/collectorargs.py

@@ -23,10 +23,22 @@
 # SOFTWARE.
 #
 ###############################################################################
-from abc import ABC
-
 from pydantic import BaseModel
 
 
-class CollectorArgs(BaseModel, ABC):
+class CollectorArgs(BaseModel):
     model_config = {"extra": "forbid", "exclude_none": True}
+
+    @classmethod
+    def build_from_model(cls, datamodel):
+        """Build analyzer args instance from data model object
+
+        Args:
+            datamodel (TDataModel): data model to use for creating analyzer args
+
+        Raises:
+            NotImplementedError: Not implemented error
+        """
+        raise NotImplementedError(
+            "Setting analyzer args from datamodel is not implemented for class: %s", cls.__name__
+        )

nodescraper/plugins/inband/process/collector_args.py

@@ -25,9 +25,20 @@
 ###############################################################################
 
 from nodescraper.models import CollectorArgs
+from nodescraper.plugins.inband.process.processdata import ProcessDataModel
 
 
 class ProcessCollectorArgs(CollectorArgs):
     top_n_process: int = 10
 
-    model_config = {"extra": "forbid"}
+    @classmethod
+    def build_from_model(cls, datamodel: ProcessDataModel) -> "ProcessCollectorArgs":
+        """build analyzer args from data model
+
+        Args:
+            datamodel (BiosDataModel): data model for plugin
+
+        Returns:
+            ProcessCollectorArgs: instance of analyzer args class
+        """
+        return cls(top_n_process=datamodel.kfd_process)

test/unit/framework/common/shared_utils.py

@@ -72,11 +72,11 @@ class TestModelArg(AnalyzerArgs):
 
     @classmethod
     def build_from_model(cls, model):
-        return cls(model_attr=int(model.some_version))
+        return cls(model_attr=int(model.foo))
 
 
 class DummyDataModel(DataModel):
-    some_version: str = None
+    foo: str = None
 
 
 class TestPluginA(PluginInterface[MockConnectionManager, None]):

test/unit/framework/test_cli_helper.py

@@ -31,11 +31,17 @@
 from types import SimpleNamespace
 
 import pytest
-from common.shared_utils import DummyDataModel
+
+# from common.shared_utils import DummyDataModel
+from conftest import DummyDataModel
 from pydantic import BaseModel
 
 from nodescraper.cli import cli
-from nodescraper.cli.helper import build_config, find_datamodel_and_result
+from nodescraper.cli.helper import (
+    build_config,
+    extract_collector_args_from_model,
+    find_datamodel_and_result,
+)
 from nodescraper.configregistry import ConfigRegistry
 from nodescraper.enums import ExecutionStatus, SystemInteractionLevel
 from nodescraper.models import PluginConfig, TaskResult
@@ -50,7 +56,7 @@ def test_generate_reference_config(plugin_registry):
             source="TestPluginA",
             message="Plugin tasks completed successfully",
             result_data=DataPluginResult(
-                system_data=DummyDataModel(some_version="17"),
+                system_data=DummyDataModel(foo="17"),
                 collection_result=TaskResult(
                     status=ExecutionStatus.OK,
                     message="BIOS: 17",
@@ -176,3 +182,59 @@ def build_from_model(cls, datamodel):
     assert isinstance(cfg, PluginConfig)
     assert set(cfg.plugins) == {parent}
     assert cfg.plugins[parent]["analysis_args"] == {}
+
+
+def test_no_collector_args_defined(caplog):
+    class NoArgsPlugin:
+        pass
+
+    caplog.set_level(logging.WARNING)
+    res = extract_collector_args_from_model(
+        NoArgsPlugin, DummyDataModel(foo=1), logging.getLogger()
+    )
+    assert res is None
+    assert "does not support reference config creation" in caplog.text
+
+
+def test_collector_args_none(caplog):
+    class NoneArgsPlugin:
+        COLLECTOR_ARGS = None
+
+    caplog.set_level(logging.WARNING)
+    res = extract_collector_args_from_model(
+        NoneArgsPlugin, DummyDataModel(foo=1), logging.getLogger()
+    )
+    assert res is None
+    assert "does not support reference config creation" in caplog.text
+
+
+def test_successful_build_from_model():
+    class MyArgs(BaseModel):
+        a: int
+
+        @classmethod
+        def build_from_model(cls, dm):
+            return cls(a=dm.foo)
+
+    class GoodPlugin:
+        COLLECTOR_ARGS = MyArgs
+
+    dm = DummyDataModel(foo=42)
+    args = extract_collector_args_from_model(GoodPlugin, dm, logging.getLogger())
+    assert isinstance(args, MyArgs)
+    assert args.a == 42
+
+
+def test_build_from_model_not_implemented(caplog):
+    class BadArgs(BaseModel):
+        @classmethod
+        def build_from_model(cls, dm):
+            raise NotImplementedError("not implemented")
+
+    class BadPlugin:
+        COLLECTOR_ARGS = BadArgs
+
+    caplog.set_level(logging.INFO)
+    res = extract_collector_args_from_model(BadPlugin, DummyDataModel(foo=7), logging.getLogger())
+    assert res is None
+    assert "not implemented" in caplog.text

test/unit/framework/test_collectorargs.py

@@ -0,0 +1,25 @@
+import pytest
+
+from nodescraper.models import CollectorArgs
+
+
+class MyArgs(CollectorArgs):
+    args_foo: int
+
+    @classmethod
+    def build_from_model(cls, datamodel):
+        return cls(args_foo=datamodel.foo)
+
+
+def test_build_from_model(dummy_data_model):
+    dummy = dummy_data_model(foo=1)
+    args = MyArgs.build_from_model(dummy)
+    assert isinstance(args, MyArgs)
+    assert args.args_foo == dummy.foo
+    dump = args.model_dump(mode="json", exclude_none=True)
+    assert dump == {"args_foo": 1}
+
+
+def test_base_build_from_model_not_implemented():
+    with pytest.raises(NotImplementedError):
+        CollectorArgs.build_from_model("anything")