Fix unit tests

Ostorlab · Dec 6, 2024 · 1ba0e76 · 1ba0e76
1 parent 4a866ce
commit 1ba0e76
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 11 deletions.
diff --git a/agent/subfinder_agent.py b/agent/subfinder_agent.py
@@ -46,11 +46,6 @@ def set_virustotal_api_key(
         if virustotal_key not in config["virustotal"]:
             config["virustotal"].append(virustotal_key)
     else:
-        # add to sources
-        if "sources" not in config:
-            config["sources"] = ["virustotal"]
-        elif "virustotal" not in config["sources"]:
-            config["sources"].append("virustotal")
         config["virustotal"] = [virustotal_key]
 
     try:

diff --git a/tests/subfinder_agent_test.py b/tests/subfinder_agent_test.py
@@ -149,7 +149,6 @@ def testSetVirusTotalApiKey_createsSectionAndAddsKeyWhenNoSectionExists() -> Non
         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"]
 
 
 def testSetVirusTotalApiKey_whenVirusTotalSectionExists_addsKeyToExistingSection() -> (
@@ -173,7 +172,6 @@ def testSetVirusTotalApiKey_whenVirusTotalSectionExists_addsKeyToExistingSection
         updated_config = yaml.load(fake_file.read_text()) or {}
         assert "virustotal" in updated_config
         assert updated_config["virustotal"] == ["example-api-key", "new_key"]
-        assert "sources" in updated_config and "virustotal" in updated_config["sources"]
 
 
 def testSetVirusTotalApiKey_whenKeyAlreadyExists_doesNotAddKeyAgain() -> None:
@@ -195,12 +193,11 @@ def testSetVirusTotalApiKey_whenKeyAlreadyExists_doesNotAddKeyAgain() -> None:
         updated_config = yaml.load(fake_file.read_text()) or {}
         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:
+def testSetVirusTotalApiKey_whenFileEmpty_addVirusTotalKey() -> None:
     """
-    Test that the function creates a `sources` section and adds the `virustotal` key
+    Test that the function adds the `virustotal` key
     when the configuration file is empty.
     """
     fake_file_path = "/fake/path/provider-config-empty.yaml"
@@ -215,4 +212,3 @@ def testSetVirusTotalApiKey_whenFileEmpty_createSourcesAndAddVirusTotalKey() ->
         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"]