diff --git a/garak/interactive.py b/garak/interactive.py index 22d5f71a2..6e6b5385b 100644 --- a/garak/interactive.py +++ b/garak/interactive.py @@ -107,7 +107,9 @@ def do_probe(self, args): logger.warning(f"Probe already set. Resetting probe to {args.probe}") print(f"Executing {args.probe}") self._cmd.probe = args.probe - elif not args.probe and not self._cmd.probe: + elif args.probe: + self._cmd.probe = args.probe + elif not self._cmd.probe: logger.warning("No probe set and no probe specified.") return None try: @@ -130,9 +132,11 @@ def do_probe(self, args): except ImportError as e: logger.error(e) print("Could not load generator from Garak generators.") + return except AttributeError as e: logger.error(e) print("Please check your generator model name.") + return evaluator = ThresholdEvaluator(self._cmd.eval_threshold) harness = garak.harnesses.probewise.ProbewiseHarness() diff --git a/tests/test_interactive.py b/tests/test_interactive.py new file mode 100644 index 000000000..e8604b2c8 --- /dev/null +++ b/tests/test_interactive.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from garak.interactive import GarakCommands + +# do_probe is wrapped by @cmd2.with_argparser; call the original function +# directly (functools.wraps preserves __wrapped__) so tests don't need a +# full cmd2.Cmd/CommandSet registration or a real garak config. +_do_probe = GarakCommands.do_probe.__wrapped__ + + +def _make_cmds(probe=None, generator=None): + cmds = GarakCommands() + cmds._cmd = SimpleNamespace( + target_type="test", + target_model="test", + probe=probe, + generator=generator, + eval_threshold=0.5, + ) + return cmds + + +def test_do_probe_sets_probe_when_none_was_previously_set(): + """A probe name passed on the first `probe ` call (no probe set yet) + must actually be stored, not silently dropped.""" + cmds = _make_cmds(probe=None) + args = SimpleNamespace(probe="mynewprobe") + + with ( + patch("garak.interactive.ThresholdEvaluator", MagicMock()), + patch("garak.harnesses.probewise.ProbewiseHarness", MagicMock()), + ): + _do_probe(cmds, args) + + assert cmds._cmd.probe == "mynewprobe" + + +def test_do_probe_does_not_crash_when_generator_fails_to_load(): + """If loading the generator raises ImportError/AttributeError, do_probe + must stop instead of falling through to use the never-assigned + `generator` local variable.""" + cmds = _make_cmds(probe=None) + args = SimpleNamespace(probe="mynewprobe") + + threshold_evaluator = MagicMock() + harness_cls = MagicMock() + with ( + patch("garak._plugins.load_plugin", side_effect=ImportError("boom")), + patch("garak.interactive.ThresholdEvaluator", threshold_evaluator), + patch("garak.harnesses.probewise.ProbewiseHarness", harness_cls), + ): + _do_probe(cmds, args) + + threshold_evaluator.assert_not_called() + harness_cls.assert_not_called() + + +def test_do_probe_does_not_crash_when_generator_name_is_invalid(): + """Same as above, for the AttributeError branch.""" + cmds = _make_cmds(probe=None) + args = SimpleNamespace(probe="mynewprobe") + + threshold_evaluator = MagicMock() + harness_cls = MagicMock() + with ( + patch("garak._plugins.load_plugin", side_effect=AttributeError("boom")), + patch("garak.interactive.ThresholdEvaluator", threshold_evaluator), + patch("garak.harnesses.probewise.ProbewiseHarness", harness_cls), + ): + _do_probe(cmds, args) + + threshold_evaluator.assert_not_called() + harness_cls.assert_not_called()