import logging from importlib import resources from typing import Any import pytest from airflow_mcp_server.parser.operation_parser import OperationDetails, OperationParser from pydantic import BaseModel logging.basicConfig(level=logging.DEBUG) logger = logging.getLogger(__name__) @pytest.fixture def spec_file(): """Get content of the v1.yaml spec file.""" with resources.files("airflow_mcp_server.resources").joinpath("v1.yaml").open("rb") as f: return f.read() @pytest.fixture def parser(spec_file) -> OperationParser: """Create OperationParser instance.""" return OperationParser(spec_path=spec_file) def test_parse_operation_basic(parser: OperationParser) -> None: """Test basic operation parsing.""" operation = parser.parse_operation("get_dags") assert isinstance(operation, OperationDetails) assert operation.operation_id == "get_dags" assert operation.path == "/dags" assert operation.method == "get" assert ( operation.description == """List DAGs in the database. `dag_id_pattern` can be set to match dags of a specific pattern """ ) assert isinstance(operation.parameters, dict) def test_parse_operation_with_no_description_but_summary(parser: OperationParser) -> None: """Test parsing operation with no description but summary.""" operation = parser.parse_operation("get_connections") assert isinstance(operation, OperationDetails) assert operation.operation_id == "get_connections" assert operation.path == "/connections" assert operation.method == "get" assert operation.description == "List connections" assert isinstance(operation.parameters, dict) def test_parse_operation_with_path_params(parser: OperationParser) -> None: """Test parsing operation with path parameters.""" operation = parser.parse_operation("get_dag") assert operation.path == "/dags/{dag_id}" assert isinstance(operation.input_model, type(BaseModel)) # Verify path parameter field exists fields = operation.input_model.__annotations__ assert "dag_id" in fields assert str in fields["dag_id"].__args__ # Check if str is in the Union types # Verify parameter is mapped correctly assert "dag_id" in operation.input_model.model_config["parameter_mapping"]["path"] def test_parse_operation_with_query_params(parser: OperationParser) -> None: """Test parsing operation with query parameters.""" operation = parser.parse_operation("get_dags") # Verify query parameter field exists fields = operation.input_model.__annotations__ assert "limit" in fields assert int in fields["limit"].__args__ # Check if int is in the Union types # Verify parameter is mapped correctly assert "limit" in operation.input_model.model_config["parameter_mapping"]["query"] def test_parse_operation_with_body_params(parser: OperationParser) -> None: """Test parsing operation with request body.""" operation = parser.parse_operation("post_dag_run") # Verify body fields exist fields = operation.input_model.__annotations__ assert "dag_run_id" in fields assert str in fields["dag_run_id"].__args__ # Check if str is in the Union types # Verify parameter is mapped correctly assert "dag_run_id" in operation.input_model.model_config["parameter_mapping"]["body"] def test_parse_operation_not_found(parser: OperationParser) -> None: """Test error handling for non-existent operation.""" with pytest.raises(ValueError, match="Operation invalid_op not found in spec"): parser.parse_operation("invalid_op") def test_extract_parameters_empty(parser: OperationParser) -> None: """Test parameter extraction with no parameters.""" params = parser.extract_parameters({}) assert isinstance(params, dict) assert "path" in params assert "query" in params assert "header" in params assert all(isinstance(v, dict) for v in params.values()) def test_map_parameter_schema_array(parser: OperationParser) -> None: """Test mapping array parameter schema.""" param: dict[str, Any] = { "name": "tags", "in": "query", "schema": {"type": "array", "items": {"type": "string"}}, } result = parser._map_parameter_schema(param) assert isinstance(result["type"], type(list)) def test_map_parameter_schema_nullable(parser: OperationParser) -> None: """Test mapping nullable parameter schema.""" param: dict[str, Any] = { "name": "test", "in": "query", "schema": {"type": "string", "nullable": True}, } result = parser._map_parameter_schema(param) # Check that str is in the Union types assert str in result["type"].__args__ assert None.__class__ in result["type"].__args__ # Check for NoneType assert not result["required"] def test_create_model_invalid_schema(parser: OperationParser) -> None: """Test error handling for invalid schema.""" with pytest.raises(ValueError, match="Schema must be an object type"): parser._create_model("Test", {"type": "string"}) def test_create_model_nested_objects(parser: OperationParser) -> None: """Test creating model with nested objects.""" schema = { "type": "object", "properties": {"nested": {"type": "object", "properties": {"field": {"type": "string"}}}}, } model = parser._create_model("Test", schema) assert issubclass(model, BaseModel) fields = model.__annotations__ assert "nested" in fields assert issubclass(fields["nested"], BaseModel) nested_fields = fields["nested"].__annotations__ assert "field" in nested_fields assert isinstance(nested_fields["field"], type(str)) def test_parse_operation_with_allof_body(parser: OperationParser) -> None: """Test parsing operation with allOf schema in request body.""" operation = parser.parse_operation("test_connection") assert isinstance(operation, OperationDetails) assert operation.operation_id == "test_connection" assert operation.path == "/connections/test" assert operation.method == "post" # Verify input model includes fields from allOf schema fields = operation.input_model.__annotations__ assert "connection_id" in fields, "Missing connection_id from ConnectionCollectionItem" assert str in fields["connection_id"].__args__, "connection_id should be a string" assert "password" in fields, "Missing password from Connection" assert str in fields["password"].__args__, "password should be a string" assert "connection_schema" in fields, "Missing schema field (aliased as connection_schema)" assert str in fields["connection_schema"].__args__, "connection_schema should be a string" # Verify parameter mapping mapping = operation.input_model.model_config["parameter_mapping"] assert "body" in mapping assert "connection_id" in mapping["body"] assert "password" in mapping["body"] assert "connection_schema" in mapping["body"] # Verify alias configuration model_fields = operation.input_model.model_fields assert "connection_schema" in model_fields assert model_fields["connection_schema"].alias == "schema", "connection_schema should alias to schema"