From 1635d0b2ddcf2fc69b7bcb3495a5579c101eee26 Mon Sep 17 00:00:00 2001 From: elyousfi Date: Fri, 6 Dec 2024 11:47:30 +0100 Subject: [PATCH] Fix unit tests --- agent/subfinder_agent.py | 4 +++- tests/subfinder_agent_test.py | 25 +++++++++++++++++++++++-- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/agent/subfinder_agent.py b/agent/subfinder_agent.py index 6ea496a..c902157 100644 --- a/agent/subfinder_agent.py +++ b/agent/subfinder_agent.py @@ -47,7 +47,9 @@ def set_virustotal_api_key( config["virustotal"].append(virustotal_key) else: # add to sources - if "sources" in config and "virustotal" not in config["sources"]: + if "sources" not in config: + config["sources"] = ["virustotal"] + elif "virustotal" not in config["sources"]: config["sources"].append("virustotal") config["virustotal"] = [virustotal_key] diff --git a/tests/subfinder_agent_test.py b/tests/subfinder_agent_test.py index ddf59e9..6867397 100644 --- a/tests/subfinder_agent_test.py +++ b/tests/subfinder_agent_test.py @@ -118,10 +118,11 @@ def testSetVirusTotalApiKey_whenWriteConfigurationFail_handleWriteError( 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) + mocker.patch("ruamel.yaml.main.YAML.dump", side_effect=IOError) sub_agent.set_virustotal_api_key( - "existing_key", str(pathlib.Path(__file__).parent / "provider-config.yaml") + "existing_key", + str(pathlib.Path(__file__).parent / "provider-config-virustotal.yaml"), ) assert "Failed to write configuration file" in caplog.text @@ -195,3 +196,23 @@ def testSetVirusTotalApiKey_whenKeyAlreadyExists_doesNotAddKeyAgain() -> None: assert "virustotal" in updated_config assert updated_config["virustotal"] == ["example-api-key"] assert "sources" in updated_config and "virustotal" in updated_config["sources"] + + +def testSetVirusTotalApiKey_whenFileEmpty_createSourcesAndAddVirusTotalKey() -> None: + """ + Test that the function creates a `sources` section and adds the `virustotal` key + when the configuration file is empty. + """ + fake_file_path = "/fake/path/provider-config-empty.yaml" + + with fake_filesystem_unittest.Patcher() as patcher: + patcher.fs.create_file(fake_file_path, contents="") + + sub_agent.set_virustotal_api_key("new_key", fake_file_path) + + yaml = ruamel.yaml.YAML(typ="safe") + fake_file = pathlib.Path(fake_file_path) + updated_config = yaml.load(fake_file.read_text()) or {} + assert "virustotal" in updated_config + assert updated_config["virustotal"] == ["new_key"] + assert "sources" in updated_config and "virustotal" in updated_config["sources"]