from io import StringIO
from unittest.mock import patch, MagicMock
from urlhaus import urlhaus
import unittest


class TestUrlhausMethods(unittest.TestCase):

    def test_main(self):
        output = {"foo": "bar"}
        with patch('sys.stdout', new=StringIO()) as mock_stdout:
            with patch('urlhaus.urlhaus.analyze', new=MagicMock(return_value=output)) as mock:
                urlhaus.main()
                expected = '{"foo": "bar"}\n'
                self.assertEqual(mock_stdout.getvalue(), expected)
                mock.assert_called_once_with('.')

    def test_buildReq(self):
        result = urlhaus.buildReq("test")
        self.assertEqual("test", result["url"])

    def test_sendReq(self):
        with patch('requests.request', new=MagicMock(return_value=MagicMock())) as mock:
            meta = {"baseUrl": "myurl"}
            response = urlhaus.sendReq(meta, "mypayload")
            mock.assert_called_once_with("POST", "myurl", data="mypayload")
            self.assertIsNotNone(response)

    def test_prepareResults_none(self):
        raw = {"query_status": "no_results"}
        results = urlhaus.prepareResults(raw)
        self.assertEqual(results["response"], raw)
        self.assertEqual(results["summary"], "no_results")
        self.assertEqual(results["status"], "info")

    def test_prepareResults_invalidUrl(self):
        raw = {"query_status": "invalid_url"}
        results = urlhaus.prepareResults(raw)
        self.assertEqual(results["response"], raw)
        self.assertEqual(results["summary"], "invalid_url")
        self.assertEqual(results["status"], "error")

    def test_prepareResults_threat(self):
        raw = {"query_status": "invalid_url"}  # This is overrided in this scenario
        raw["threat"] = "bad_actor"
        results = urlhaus.prepareResults(raw)
        self.assertEqual(results["response"], raw)
        self.assertEqual(results["summary"], "bad_actor")
        self.assertEqual(results["status"], "danger")

    def test_analyze(self):
        output = {"threat": "malware_download"}
        artifactInput = '{"value":"foo","artifactType":"url"}'
        with patch('urlhaus.urlhaus.sendReq', new=MagicMock(return_value=output)) as mock:
            results = urlhaus.analyze(artifactInput)
            self.assertEqual(results["summary"], "malware_download")
            mock.assert_called_once()