import unittest

from fgai.llm import _json_object_from_text


class LlmTests(unittest.TestCase):
    def test_json_object_from_text_accepts_markdown_wrapped_json(self):
        payload = _json_object_from_text(
            '```json\n{"profiles":[{"stream_id":"windows","entity_fields":["username"]}]}\n```'
        )

        self.assertEqual(payload["profiles"][0]["stream_id"], "windows")

    def test_json_object_from_text_extracts_object_from_extra_text(self):
        payload = _json_object_from_text(
            'Here is the profile:\n{"profiles":[{"stream_id":"firewall"}]}\nDone.'
        )

        self.assertEqual(payload["profiles"][0]["stream_id"], "firewall")


if __name__ == "__main__":
    unittest.main()