test_agent.py

import unittest
import os
from unittest.mock import patch, MagicMock
from agent import (
    build_graph,
    multiply,
    add,
    subtract,
    divide,
    modulus,
    power,
    square_root,
    save_and_read_file,
    download_file_from_url,
    extract_text_from_image,
    analyze_image,
    transform_image,
    draw_on_image,
    generate_simple_image,
    combine_images,
    analyze_csv_file,
    analyze_excel_file,
    execute_code_multilang,
    web_search,
    wiki_search,
    arxiv_search,
    vector_store,
)

class TestAgent(unittest.TestCase):
    def test_multiply(self):
        response = multiply.invoke({"a": 6, "b": 7})
        self.assertTrue(response["status"])
        self.assertEqual(response["data"], 42)

    def test_add(self):
        response = add.invoke({"a": 5, "b": 3})
        self.assertTrue(response["status"])
        self.assertEqual(response["data"], 8)

    def test_llm(self):
        graph = build_graph()
        response = graph.invoke({"messages": ["what is 1 + 1"]})
        self.assertIsNotNone(response)
        print(response)


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