242 lines
8.9 KiB
Python
242 lines
8.9 KiB
Python
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()
|