import unittest
from unittest.mock import patch, MagicMock, mock_open
import json
import asyncio
from jsonschema import ValidationError
with patch.dict('os.environ', {
    "REPO_NAME": "test_repo",
    "BRANCH_NAME": "test_branch",
    "VERSION": "test_version",
    "NAMESPACE": "test_namespace",
    "FLOWX_ENGINE_ADDRESS": "test_address"
}):
    from block_wrapper import block_main_activity, validate_input, validate_output


class TestBlockWrapper(unittest.TestCase):

    def setUp(self):
        # Mock schemas to use for testing
        self.mock_request_schema = {
            "type": "object",
            "properties": {
                "a": {"type": "integer"},
                "b": {"type": "integer"}
            },
            "required": ["a", "b"]
        }
        self.mock_response_schema = {
            "type": "object",
            "properties": {
                "sum": {"type": "integer"}
            },
            "required": ["sum"]
        }

        # Mock the contents of request_schema.json and response_schema.json using different patchers
        self.mock_open_request = mock_open(read_data=json.dumps(self.mock_request_schema))
        self.mock_open_response = mock_open(read_data=json.dumps(self.mock_response_schema))

        # Mock load_block_main to return a mock main function
        self.load_block_main_patcher = patch("block_wrapper.load_block_main", return_value=MagicMock(return_value={"sum": 3}))
        self.mock_load_block_main = self.load_block_main_patcher.start()

    def tearDown(self):
        # Stop all patches
        self.load_block_main_patcher.stop()

    @patch("block_wrapper.load_schema")
    def test_validate_input_success(self, mock_load_schema):
        # Set up load_schema to return request schema for validate_input
        mock_load_schema.return_value = self.mock_request_schema
        input_data = {"a": 1, "b": 2}
        validate_input(input_data)  # Should pass without errors

    @patch("block_wrapper.load_schema")
    def test_validate_input_failure(self, mock_load_schema):
        # Set up load_schema to return request schema for validate_input
        mock_load_schema.return_value = self.mock_request_schema
        input_data = {"a": 1}  # Missing 'b'
        with self.assertRaises(ValueError):
            validate_input(input_data)

    @patch("block_wrapper.load_schema")
    def test_validate_output_success(self, mock_load_schema):
        # Set up load_schema to return response schema for validate_output
        mock_load_schema.return_value = self.mock_response_schema
        output_data = {"sum": 3}
        validate_output(output_data)  # Should pass without errors

    @patch("block_wrapper.load_schema")
    def test_validate_output_failure(self, mock_load_schema):
        # Set up load_schema to return response schema for validate_output
        mock_load_schema.return_value = self.mock_response_schema
        output_data = {}  # Missing 'sum'
        with self.assertRaises(ValueError):
            validate_output(output_data)

    @patch("block_wrapper.load_schema")
    async def test_block_main_activity_success(self, mock_load_schema):
        # Set up load_schema to return request and response schemas in order
        mock_load_schema.side_effect = [self.mock_request_schema, self.mock_response_schema]
        input_data = {"a": 1, "b": 2}
        result = await block_main_activity(input_data)
        self.assertEqual(result, {"sum": 3})

    @patch("block_wrapper.load_schema")
    async def test_block_main_activity_failure(self, mock_load_schema):
        # Set up load_schema to return request and response schemas in order
        mock_load_schema.side_effect = [self.mock_request_schema, self.mock_response_schema]
        # Cause an exception in main function
        self.mock_load_block_main.side_effect = Exception("Unexpected error")
        input_data = {"a": 1, "b": 2}
        with self.assertRaises(RuntimeError):
            await block_main_activity(input_data)

    @patch("block_wrapper.load_schema")
    async def test_block_main_activity_input_validation_failure(self, mock_load_schema):
        # Mock validate_input to raise ValidationError
        with patch("block_wrapper.validate_input", side_effect=ValidationError("Invalid input")):
            input_data = {"a": 1}  # Missing 'b'
            with self.assertRaises(ValueError):
                await block_main_activity(input_data)

    @patch("block_wrapper.load_schema")
    async def test_block_main_activity_output_validation_failure(self, mock_load_schema):
        # Mock validate_output to raise ValidationError
        with patch("block_wrapper.validate_output", side_effect=ValidationError("Invalid output")):
            input_data = {"a": 1, "b": 2}
            with self.assertRaises(ValueError):
                await block_main_activity(input_data)

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