diff --git a/tests/unittest/try_fix_yaml.py b/tests/unittest/try_fix_yaml.py new file mode 100644 index 00000000..fe37bb1c --- /dev/null +++ b/tests/unittest/try_fix_yaml.py @@ -0,0 +1,31 @@ + +# Generated by CodiumAI +from pr_agent.algo.utils import try_fix_yaml + + +import pytest + +class TestTryFixYaml: + + # The function successfully parses a valid YAML string. + def test_valid_yaml(self): + review_text = "key: value\n" + expected_output = {"key": "value"} + assert try_fix_yaml(review_text) == expected_output + + # The function adds '|-' to 'relevant line:' if it is not already present and successfully parses the YAML string. + def test_add_relevant_line(self): + review_text = "relevant line: value: 3\n" + expected_output = {"relevant line": "value: 3"} + assert try_fix_yaml(review_text) == expected_output + + # The function removes the last line(s) of the YAML string and successfully parses the YAML string. + def test_remove_last_line(self): + review_text = "key: value\nextra invalid line\n" + expected_output = {"key": "value"} + assert try_fix_yaml(review_text) == expected_output + + # The YAML string is empty. + def test_empty_yaml_fixed(self): + review_text = "" + assert try_fix_yaml(review_text) is None \ No newline at end of file