update tests

Author: paulinek13

tests/unit/common/test_helper_functions.py

@@ -1,6 +1,8 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+from unittest.mock import patch
+
 from pyrit.common.utils import combine_dict, select_word_indices
 
 
@@ -36,5 +38,12 @@ def test_combine_dict_same_keys():
 
 def test_word_indices_selection():
     assert select_word_indices(words=["word1", "word2", "word3"], mode="all") == [0, 1, 2]
+    assert select_word_indices(words=["word1", "word2", "word3"], mode="custom", indices=[0, 2]) == [0, 2]
     assert select_word_indices(words=["word1", "word2", "pyrit", "word4"], mode="keywords", keywords=["pyrit"]) == [2]
     assert select_word_indices(words=["word1", "word2", "pyrit", "word4"], mode="regex", regex=r"word\d") == [0, 1, 3]
+
+    with patch("random.sample", return_value=[0, 2]):
+        result = select_word_indices(words=["word1", "word2", "word3", "word4"], mode="random", sample_ratio=0.5)
+    assert result == [0, 2]
+
+    assert select_word_indices(words=["word1", "word2"], mode="invalid_mode") == [0, 1]