From 5d621646e6b19813cd2b7316f0fdf59a6faab606 Mon Sep 17 00:00:00 2001
From: pirahnasa <saad@ostorlab.dev>
Date: Wed, 4 Dec 2024 18:21:36 +0100
Subject: [PATCH] fix codecov

---
 agent/subfinder_agent.py      |  8 +++---
 tests/conftest.py             | 27 +++++++++++++++++++
 tests/provider-config.yaml    |  0
 tests/subfinder_agent_test.py | 51 +++++++++++++++++++++++++++++++++++
 4 files changed, 82 insertions(+), 4 deletions(-)
 create mode 100644 tests/provider-config.yaml

diff --git a/agent/subfinder_agent.py b/agent/subfinder_agent.py
index 146f225..d156331 100644
--- a/agent/subfinder_agent.py
+++ b/agent/subfinder_agent.py
@@ -26,7 +26,7 @@
 CONFIG_PATH = "/root/.config/subfinder/provider-config.yaml"
 
 
-def _update_provider_config(
+def update_provider_config(
     virustotal_key: str,
     config_path: str = CONFIG_PATH,
 ) -> None:
@@ -38,8 +38,8 @@ def _update_provider_config(
         with open(config_path, "r") as config_file:
             config = yaml.load(config_file) or {}
     except FileNotFoundError:
-        logger.warning("Configuration file not found. Creating a new one.")
-
+        logger.error("Configuration file not found. Creating a new one.")
+        return
     # Update the 'virustotal' section
     if "virustotal" in config:
         if virustotal_key not in config["virustotal"]:
@@ -69,7 +69,7 @@ def __init__(
         virustotal_key = self.args.get("virustotal_key")
         if virustotal_key is not None:
             logger.info("Updating configuration with VirusTotal API key.")
-            _update_provider_config(virustotal_key)
+            update_provider_config(virustotal_key)
 
         agent_persist_mixin.AgentPersistMixin.__init__(self, agent_settings)
 
diff --git a/tests/conftest.py b/tests/conftest.py
index ff3e9b2..46c62e0 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -49,3 +49,30 @@ def fixture_subfinder_agent_max_subdomains():
 
         agent = subfinder_agent.SubfinderAgent(definition, settings)
         return agent
+
+
+@pytest.fixture
+def subfinder_definition() -> agent_definitions.AgentDefinition:
+    with (pathlib.Path(__file__).parent.parent / "ostorlab.yaml").open() as yaml_o:
+        definition = agent_definitions.AgentDefinition.from_yaml(yaml_o)
+        definition.args = [
+            {
+                "name": "virustotal_key",
+                "value": "Justrandomvalue",
+                "type": "string",
+            },
+        ]
+        return definition
+
+
+@pytest.fixture
+def subfinder_settings() -> runtime_definitions.AgentSettings:
+    settings = runtime_definitions.AgentSettings(
+        key="agent/ostorlab/subfinder",
+        bus_url="NA",
+        bus_exchange_topic="NA",
+        args=[],
+        healthcheck_port=random.randint(5000, 6000),
+        redis_url="redis://guest:guest@localhost:6379",
+    )
+    return settings
diff --git a/tests/provider-config.yaml b/tests/provider-config.yaml
new file mode 100644
index 0000000..e69de29
diff --git a/tests/subfinder_agent_test.py b/tests/subfinder_agent_test.py
index a360f2d..4eb08ba 100644
--- a/tests/subfinder_agent_test.py
+++ b/tests/subfinder_agent_test.py
@@ -1,6 +1,14 @@
 """Unittests for the Subfinder Agent."""
 
+import pathlib
+
+import pytest
+from pytest_mock import plugin
+
 from ostorlab.agent.message import message
+from agent import subfinder_agent as sub_agent
+from ostorlab.agent import definitions as agent_definitions
+from ostorlab.runtimes import definitions as runtime_definitions
 
 
 def testAgentSubfinder_whenFindsSubDomains_emitsBackFindings(
@@ -72,3 +80,46 @@ def testAgentSubfinder_whenMaxSubDomainsSet_emitsBackFindings(
     assert agent_mock[0].selector == "v3.asset.domain_name", (
         agent_mock[0].data["name"] == "subdomain1"
     )
+
+
+def testAgentSubfinder_whenVirustotalKeyPassed_emitsBackFindings(
+    subfinder_definition: agent_definitions.AgentDefinition,
+    subfinder_settings: runtime_definitions.AgentSettings,
+    mocker: plugin.MockerFixture,
+) -> None:
+    """
+    Test that the Subfinder agent correctly updates the provider configuration
+    with the VirusTotal key.
+    """
+    mocker_update_provider_config = mocker.patch(
+        "agent.subfinder_agent.update_provider_config"
+    )
+
+    sub_agent.SubfinderAgent(subfinder_definition, subfinder_settings)
+
+    assert mocker_update_provider_config.called is True
+    assert mocker_update_provider_config.call_args[0][0] == "Justrandomvalue"
+
+
+def testUpdateConfigurationFile_whenConfNotFound_handelFileNotFoundError(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that the provider configuration correctly handles a missing configuration file."""
+
+    sub_agent.update_provider_config("existing_key", "test_not.yaml")
+
+    assert "Configuration file not found. Creating a new one." in caplog.text
+
+
+def testupdateconfigupdate_whenWriteConfigurationFail_handelWriteErro(
+    mocker: plugin.MockerFixture,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that the provider configuration handles a write error correctly and logs the failure."""
+    mocker.patch("ruamel.yaml.main.YAML.dump", side_effect=FileNotFoundError)
+
+    sub_agent.update_provider_config(
+        "existing_key", str(pathlib.Path(__file__).parent / "provider-config.yaml")
+    )
+
+    assert "Failed to write configuration file" in caplog.text