import unittest from unittest.mock import patch, MagicMock, mock_open import json import asyncio from jsonschema import ValidationError import os 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, construct_sql, get_connection_id class TestBlockWrapper(unittest.TestCase): def setUp(self): # Mock schemas to use for testing self.mock_request_schema = { "type": "object", "properties": { "salary": {"type": "number"}, "department": {"type": "string"} }, "required": ["salary", "department"] } self.mock_response_schema = { "type": "object", "$schema": "http://json-schema.org/draft-07/schema", "properties": { "id": { "type": "integer" }, "first_name": { "type": "string" }, "last_name": { "type": "string" }, "email": { "type": "string", "format": "email" }, "phone_number": { "type": "string" }, "hire_date": { "type": "string", "format": "date-time" }, "job_title": { "type": "string" }, "salary": { "type": "number" }, "department": { "type": "string" } }, "required": ["id","first_name","last_name","email","phone_number","hire_date","job_title","salary","department"] } self.mock_config_schema = [ { "namespace": "staging", "connectionId": "8d7341b4-53a5-41b8-8c9d-5133fafb5d7b" }, { "namespace": "production", "connectionId": "4b1437d8-53a5-41b8-8c9d-5133fafbtyuu" } ] self.mock_main_sql = "SELECT * FROM public.employee WHERE salary=$salary and department=$department;" # 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)) self.mock_open_config = mock_open(read_data=json.dumps(self.mock_config_schema)) self.mock_open_main_sql = mock_open(read_data=self.mock_main_sql) self.open_main_sql_patcher = patch("builtins.open", self.mock_open_main_sql) self.open_main_sql_patcher.start() # Mock execute_sqlpad_query to return a known result self.load_block_main_patcher = patch( "block_wrapper.execute_sqlpad_query", return_value=MagicMock( return_value={ "id": 4, "first_name": "Bob", "last_name": "Brown", "email": "bob.brown@example.com", "phone_number": "444-222-1111", "hire_date": "2020-07-25", "job_title": "Marketing Specialist", "salary": 60000.00, "department": "Marketing" } ) ) self.mock_load_block_main = self.load_block_main_patcher.start() def tearDown(self): # Stop all patches self.load_block_main_patcher.stop() self.open_main_sql_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 = {"salary": 20000.0, "department": "Marketing"} 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 = {"salary": 20000.00} # Missing 'department' 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 = { "id": 4, "first_name": "Bob", "last_name": "Brown", "email": "bob.brown@example.com", "phone_number": "444-222-1111", "hire_date": "2020-07-25", "job_title": "Marketing Specialist", "salary": 60000.00, "department": "Marketing" } 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 # Missing department output_data = { "id": 4, "first_name": "Bob", "last_name": "Brown", "email": "bob.brown@example.com", "phone_number": "444-222-1111", "hire_date": "2020-07-25", "job_title": "Marketing Specialist", "salary": 60000.00 } 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 = {"salary": 20000.0, "department": "Marketing"} result = await block_main_activity(input_data) self.assertEqual( result, { "id": 4, "first_name": "Bob", "last_name": "Brown", "email": "bob.brown@example.com", "phone_number": "444-222-1111", "hire_date": "2020-07-25", "job_title": "Marketing Specialist", "salary": 60000.00, "department": "Marketing" } ) @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 = {"salary": 20000.0, "department": "Marketing"} with self.assertRaises(RuntimeError): await block_main_activity(input_data) @patch("block_wrapper.validate_input", side_effect=ValidationError("Invalid input")) async def test_block_main_activity_input_validation_failure(self, mock_validate): input_data = {"salary": 20000.00} # Missing 'department' with self.assertRaises(ValueError): await block_main_activity(input_data) @patch("block_wrapper.validate_output", side_effect=ValidationError("Invalid output")) async def test_block_main_activity_output_validation_failure(self, mock_validate): input_data = {"salary": 20000.0, "department": "Marketing"} with self.assertRaises(ValueError): await block_main_activity(input_data) @patch.dict(os.environ, {"NAMESPACE": "staging"}) @patch("block_wrapper.load_schema") def test_get_connection_id_staging(self, mock_load_schema): """Test fetching connectionId for 'staging' namespace""" mock_load_schema.return_value = self.mock_config_schema connection_id = get_connection_id(os.environ["NAMESPACE"]) self.assertEqual(connection_id, "8d7341b4-53a5-41b8-8c9d-5133fafb5d7b") @patch.dict(os.environ, {"NAMESPACE": "production"}) @patch("block_wrapper.load_schema") def test_get_connection_id_production(self, mock_load_schema): """Test fetching connectionId for 'production' namespace""" mock_load_schema.return_value = self.mock_config_schema connection_id = get_connection_id(os.environ["NAMESPACE"]) self.assertEqual(connection_id, "4b1437d8-53a5-41b8-8c9d-5133fafbtyuu") @patch("block_wrapper.load_schema") def test_get_connection_id_invalid_namespace(self, mock_load_schema): """Test handling of invalid namespace""" mock_load_schema.return_value = self.mock_config_schema with self.assertRaises(ValueError) as context: get_connection_id("development") self.assertIn("Namespace 'development' not found", str(context.exception)) @patch("block_wrapper.load_schema") def test_valid_sql_replacement(self, mock_load_schema): mock_load_schema.return_value = self.mock_main_sql input_data = {"salary": 20000.0, "department": "Marketing"} expected_sql = "SELECT * FROM public.employee WHERE salary=20000.0 and department='Marketing';" result = construct_sql(input_data) self.assertEqual(result, expected_sql) if __name__ == "__main__": unittest.main()