24 lines
726 B
Python
24 lines
726 B
Python
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()
|