def testGetKBDependencies(self): """Test that KB dependencies are calculated correctly.""" artifact_registry.REGISTRY.ClearSources() try: test_artifacts_file = os.path.join( config_lib.CONFIG["Test.data_dir"], "artifacts", "test_artifacts.json") artifact_registry.REGISTRY.AddFileSource(test_artifacts_file) with test_lib.ConfigOverrider({ "Artifacts.knowledge_base": [ "DepsParent", "DepsDesktop", "DepsHomedir", "DepsWindir", "DepsWindirRegex", "DepsControlSet", "FakeArtifact" ], "Artifacts.knowledge_base_additions": ["DepsHomedir2"], "Artifacts.knowledge_base_skip": ["DepsWindir"], "Artifacts.knowledge_base_heavyweight": ["FakeArtifact"] }): args = artifact.KnowledgeBaseInitializationArgs( lightweight=True) kb_init = artifact.KnowledgeBaseInitializationFlow( None, token=self.token) kb_init.args = args kb_init.state.Register("all_deps", set()) kb_init.state.Register("awaiting_deps_artifacts", []) kb_init.state.Register("knowledge_base", rdf_client.KnowledgeBase(os="Windows")) no_deps = kb_init.GetFirstFlowsForCollection() self.assertItemsEqual(no_deps, ["DepsControlSet", "DepsHomedir2"]) self.assertItemsEqual(kb_init.state.all_deps, [ "users.homedir", "users.desktop", "users.username", "environ_windir", "current_control_set" ]) self.assertItemsEqual(kb_init.state.awaiting_deps_artifacts, [ "DepsParent", "DepsDesktop", "DepsHomedir", "DepsWindirRegex" ]) finally: artifact.ArtifactLoader().RunOnce()
def testGetKBDependencies(self): """Test that KB dependencies are calculated correctly.""" self.SetupWindowsMocks() with utils.Stubber(artifact_registry.ArtifactRegistry, "artifacts", {}): test_artifacts_file = os.path.join( config_lib.CONFIG["Test.data_dir"], "test_artifacts.json") artifact_lib.LoadArtifactsFromFiles([test_artifacts_file]) config_lib.CONFIG.Set("Artifacts.knowledge_base", [ "DepsParent", "DepsDesktop", "DepsHomedir", "DepsWindir", "DepsWindirRegex", "DepsControlSet", "FakeArtifact" ]) config_lib.CONFIG.Set("Artifacts.knowledge_base_additions", ["DepsHomedir2"]) config_lib.CONFIG.Set("Artifacts.knowledge_base_skip", ["DepsWindir"]) config_lib.CONFIG.Set("Artifacts.knowledge_base_heavyweight", ["FakeArtifact"]) args = artifact.KnowledgeBaseInitializationArgs(lightweight=True) kb_init = artifact.KnowledgeBaseInitializationFlow( None, token=self.token) kb_init.args = args kb_init.state.Register("all_deps", set()) kb_init.state.Register("awaiting_deps_artifacts", []) kb_init.state.Register("knowledge_base", rdf_client.KnowledgeBase(os="Windows")) no_deps = kb_init.GetFirstFlowsForCollection() self.assertItemsEqual(no_deps, ["DepsControlSet", "DepsHomedir2"]) self.assertItemsEqual(kb_init.state.all_deps, [ "users.homedir", "users.desktop", "users.username", "environ_windir", "current_control_set" ]) self.assertItemsEqual(kb_init.state.awaiting_deps_artifacts, [ "DepsParent", "DepsDesktop", "DepsHomedir", "DepsWindirRegex" ])