# ------------------------------------ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # ------------------------------------ import pytest import pandas as pd import os import pathlib from typing import Dict, Any from azure.ai.evaluation._evaluate._evaluate_aoai import ( _generate_data_source_config, _get_data_source, _build_schema_tree_from_paths, WRAPPER_KEY, ) def _get_file(name): """Get the file from the unittest data folder.""" data_path = os.path.join(pathlib.Path(__file__).parent.resolve(), "data") return os.path.join(data_path, name) @pytest.fixture def flat_test_data(): """Fixture for flat structure test data.""" return pd.DataFrame( [ { "query": "What is the capital of France?", "response": "Paris is the capital of France.", "ground_truth": "Paris", }, {"query": "What is 2+2?", "response": "The answer is 4.", "ground_truth": "4"}, { "query": "Who wrote Hamlet?", "response": "William Shakespeare wrote Hamlet.", "ground_truth": "Shakespeare", }, ] ) @pytest.fixture def nested_test_data(): """Fixture for nested structure test data.""" return pd.DataFrame( [ { "item.query": "What security policies exist?", "item.context.company.policy.security.passwords.rotation_days": "90", "item.context.company.policy.security.network.vpn.required": "true", "item.response": "Password rotation is required every 90 days.", "item.ground_truth": "Security policies include password rotation.", }, { "item.query": "What are the database settings?", "item.context.company.infrastructure.database.host": "db.example.com", "item.context.company.infrastructure.database.port": "5432", "item.response": "The database is PostgreSQL.", "item.ground_truth": "PostgreSQL database", }, ] ) @pytest.fixture def flat_test_data_file(): """Fixture for flat test data file path.""" return _get_file("flat_test_data.jsonl") @pytest.fixture def nested_test_data_file(): """Fixture for nested test data file path.""" return _get_file("nested_test_data.jsonl") @pytest.fixture def wrapped_flat_test_data_file(): """Fixture for wrapped flat test data file path.""" return _get_file("wrapped_flat_test_data.jsonl") @pytest.mark.unittest class TestBuildSchemaTreeFromPaths: """Test suite for the _build_schema_tree_from_paths helper function.""" def test_single_level_paths(self): """Test building schema with single-level paths.""" paths = ["query", "response", "ground_truth"] schema = _build_schema_tree_from_paths(paths, force_leaf_type="string") assert schema["type"] == "object" assert "properties" in schema assert "required" in schema assert set(schema["properties"].keys()) == {"query", "response", "ground_truth"} assert all(prop["type"] == "string" for prop in schema["properties"].values()) assert set(schema["required"]) == {"query", "response", "ground_truth"} def test_nested_paths(self): """Test building schema with nested paths.""" paths = [ "context.company.policy.security.passwords.rotation_days", "context.company.policy.security.network.vpn.required", "query", ] schema = _build_schema_tree_from_paths(paths, force_leaf_type="string") assert schema["type"] == "object" assert "context" in schema["properties"] assert schema["properties"]["context"]["type"] == "object" # Navigate nested structure company = schema["properties"]["context"]["properties"]["company"] assert company["type"] == "object" policy = company["properties"]["policy"] assert policy["type"] == "object" security = policy["properties"]["security"] assert security["type"] == "object" # Check leaf nodes passwords = security["properties"]["passwords"] assert passwords["properties"]["rotation_days"]["type"] == "string" network = security["properties"]["network"] assert network["properties"]["vpn"]["properties"]["required"]["type"] == "string" # Check required arrays exist at each level assert "required" in schema assert "required" in schema["properties"]["context"] def test_empty_paths(self): """Test building schema with empty paths list.""" paths = [] schema = _build_schema_tree_from_paths(paths, force_leaf_type="object") assert schema["type"] == "object" def test_mixed_depth_paths(self): """Test building schema with paths of different depths.""" paths = ["simple_field", "nested.field.deep", "nested.field.shallow", "another.path"] schema = _build_schema_tree_from_paths(paths, force_leaf_type="string") assert "simple_field" in schema["properties"] assert schema["properties"]["simple_field"]["type"] == "string" assert "nested" in schema["properties"] nested = schema["properties"]["nested"] assert nested["type"] == "object" assert "field" in nested["properties"] @pytest.mark.unittest class TestGenerateDataSourceConfig: """Test suite for the _generate_data_source_config function.""" def test_flat_column_mapping(self, flat_test_data): """Test generating data source config with flat column mappings.""" column_mapping = { "query": "${data.query}", "response": "${data.response}", "ground_truth": "${data.ground_truth}", } config = _generate_data_source_config(flat_test_data, column_mapping) assert config["type"] == "custom" assert "item_schema" in config assert config["item_schema"]["type"] == "object" properties = config["item_schema"]["properties"] assert "query" in properties assert "response" in properties assert "ground_truth" in properties # All should be strings in flat mode assert properties["query"]["type"] == "string" assert properties["response"]["type"] == "string" assert properties["ground_truth"]["type"] == "string" def test_nested_column_mapping_with_wrapper(self, nested_test_data): """Test generating data source config with nested paths under wrapper.""" column_mapping = { "query": "${data.item.query}", "passwords_rotation": "${data.item.context.company.policy.security.passwords.rotation_days}", "vpn_required": "${data.item.context.company.policy.security.network.vpn.required}", "response": "${data.item.response}", } config = _generate_data_source_config(nested_test_data, column_mapping) assert config["type"] == "custom" assert "item_schema" in config schema = config["item_schema"] # Should be nested object since paths contain dots assert schema["type"] == "object" # The wrapper should be stripped, so we should see inner structure assert "query" in schema["properties"] assert "response" in schema["properties"] assert "context" in schema["properties"] # Verify nested structure context = schema["properties"]["context"] assert context["type"] == "object" assert "company" in context["properties"] def test_nested_column_mapping_without_wrapper(self, nested_test_data): """Test generating data source config with nested paths not using standard wrapper.""" column_mapping = { "query": "${data.custom.query}", "field": "${data.custom.nested.field}", } config = _generate_data_source_config(nested_test_data, column_mapping) assert config["type"] == "custom" assert "item_schema" in config schema = config["item_schema"] # Should be nested assert schema["type"] == "object" # Without wrapper stripping, should see 'custom' at top level assert "custom" in schema["properties"] def test_mixed_data_and_run_outputs(self, flat_test_data): """Test column mapping with both data and run.outputs references.""" column_mapping = { "query": "${data.query}", "response": "${run.outputs.response}", "ground_truth": "${data.ground_truth}", } config = _generate_data_source_config(flat_test_data, column_mapping) # Only data.* paths should be in schema properties = config["item_schema"]["properties"] assert "query" in properties assert "ground_truth" in properties # run.outputs.response shouldn't create a schema property directly def test_empty_column_mapping(self, flat_test_data): """Test with empty column mapping.""" column_mapping = {} config = _generate_data_source_config(flat_test_data, column_mapping) # Should return flat schema with no properties assert config["type"] == "custom" assert config["item_schema"]["type"] == "object" assert config["item_schema"]["properties"] == {} def test_no_data_references(self, flat_test_data): """Test column mapping with no ${data.*} references.""" column_mapping = {"response": "${run.outputs.response}", "result": "${run.outputs.result}"} config = _generate_data_source_config(flat_test_data, column_mapping) # Should return flat schema since no data paths referenced assert config["type"] == "custom" assert "response" in config["item_schema"]["properties"] assert "result" in config["item_schema"]["properties"] def test_single_nested_path(self, flat_test_data): """Test with a single nested path to ensure nested mode activates.""" column_mapping = {"nested_field": "${data.item.context.field}"} config = _generate_data_source_config(flat_test_data, column_mapping) # Should generate nested schema assert config["type"] == "custom" schema = config["item_schema"] assert schema["type"] == "object" # After wrapper stripping, should see context assert "context" in schema["properties"] @pytest.mark.unittest class TestGetDataSource: """Test suite for the _get_data_source function.""" def test_flat_data_source_generation(self, flat_test_data): """Test generating data source from flat data.""" column_mapping = { "query": "${data.query}", "response": "${data.response}", "ground_truth": "${data.ground_truth}", } data_source = _get_data_source(flat_test_data, column_mapping) assert data_source["type"] == "jsonl" assert "source" in data_source assert data_source["source"]["type"] == "file_content" content = data_source["source"]["content"] assert len(content) == 3 # Each item should be wrapped for item in content: assert WRAPPER_KEY in item assert "query" in item[WRAPPER_KEY] assert "response" in item[WRAPPER_KEY] assert "ground_truth" in item[WRAPPER_KEY] def test_nested_data_source_generation(self, nested_test_data): """Test generating data source from nested data.""" column_mapping = { "query": "${data.item.query}", "rotation_days": "${data.item.context.company.policy.security.passwords.rotation_days}", "vpn_required": "${data.item.context.company.policy.security.network.vpn.required}", "response": "${data.item.response}", } data_source = _get_data_source(nested_test_data, column_mapping) assert data_source["type"] == "jsonl" content = data_source["source"]["content"] assert len(content) == 2 # Verify nested structure is built correctly first_item = content[0][WRAPPER_KEY] assert "query" in first_item assert "context" in first_item assert "company" in first_item["context"] assert "policy" in first_item["context"]["company"] # Check leaf values passwords = first_item["context"]["company"]["policy"]["security"]["passwords"] assert passwords["rotation_days"] == "90" vpn = first_item["context"]["company"]["policy"]["security"]["network"]["vpn"] assert vpn["required"] == "true" def test_data_source_with_run_outputs(self, flat_test_data): """Test data source generation with run.outputs mappings.""" # Add __outputs column to simulate target function output flat_test_data["__outputs.model_response"] = [ "Generated response 1", "Generated response 2", "Generated response 3", ] column_mapping = { "query": "${data.query}", "response": "${run.outputs.model_response}", "ground_truth": "${data.ground_truth}", } data_source = _get_data_source(flat_test_data, column_mapping) content = data_source["source"]["content"] # run.outputs should be mapped with just leaf name for i, item in enumerate(content): assert "model_response" in item[WRAPPER_KEY] assert item[WRAPPER_KEY]["model_response"] == f"Generated response {i+1}" def test_data_source_with_unmapped_columns(self, flat_test_data): """Test that unmapped columns are included in output.""" # Add extra column not in mapping flat_test_data["extra_field"] = ["extra1", "extra2", "extra3"] column_mapping = {"query": "${data.query}", "response": "${data.response}"} data_source = _get_data_source(flat_test_data, column_mapping) content = data_source["source"]["content"] # Unmapped columns should appear directly in item for i, item in enumerate(content): assert "extra_field" in item[WRAPPER_KEY] assert "ground_truth" in item[WRAPPER_KEY] # Also unmapped def test_data_source_with_none_values(self, flat_test_data): """Test data source generation handles None values correctly.""" flat_test_data.loc[1, "response"] = None column_mapping = { "query": "${data.query}", "response": "${data.response}", "ground_truth": "${data.ground_truth}", } data_source = _get_data_source(flat_test_data, column_mapping) content = data_source["source"]["content"] # None should be converted to empty string assert content[1][WRAPPER_KEY]["response"] == "" def test_data_source_with_numeric_values(self, flat_test_data): """Test data source generation converts numeric values to strings.""" flat_test_data["score"] = [95, 87, 92] flat_test_data["confidence"] = [0.95, 0.87, 0.92] column_mapping = {"query": "${data.query}", "score": "${data.score}", "confidence": "${data.confidence}"} data_source = _get_data_source(flat_test_data, column_mapping) content = data_source["source"]["content"] # Numeric values should be converted to strings assert content[0][WRAPPER_KEY]["score"] == "95" assert content[0][WRAPPER_KEY]["confidence"] == "0.95" assert isinstance(content[0][WRAPPER_KEY]["score"], str) assert isinstance(content[0][WRAPPER_KEY]["confidence"], str) def test_empty_dataframe(self): """Test data source generation with empty dataframe.""" empty_df = pd.DataFrame() column_mapping = {"query": "${data.query}"} data_source = _get_data_source(empty_df, column_mapping) assert data_source["type"] == "jsonl" assert len(data_source["source"]["content"]) == 0 def test_complex_nested_structure(self): """Test with complex multi-level nested structure.""" df = pd.DataFrame( [ { "item.a.b.c.d": "deep_value", "item.a.b.x": "mid_value", "item.a.y": "shallow_value", "item.z": "top_value", } ] ) column_mapping = { "deep": "${data.item.a.b.c.d}", "mid": "${data.item.a.b.x}", "shallow": "${data.item.a.y}", "top": "${data.item.z}", } data_source = _get_data_source(df, column_mapping) content = data_source["source"]["content"] item = content[0][WRAPPER_KEY] # Verify nested structure assert item["a"]["b"]["c"]["d"] == "deep_value" assert item["a"]["b"]["x"] == "mid_value" assert item["a"]["y"] == "shallow_value" assert item["z"] == "top_value" def test_data_source_preserves_row_order(self, flat_test_data): """Test that data source preserves the order of rows.""" column_mapping = {"query": "${data.query}", "response": "${data.response}"} data_source = _get_data_source(flat_test_data, column_mapping) content = data_source["source"]["content"] # Verify order matches input assert content[0][WRAPPER_KEY]["query"] == flat_test_data.iloc[0]["query"] assert content[1][WRAPPER_KEY]["query"] == flat_test_data.iloc[1]["query"] assert content[2][WRAPPER_KEY]["query"] == flat_test_data.iloc[2]["query"] @pytest.mark.unittest class TestDataSourceConfigIntegration: """Integration tests for schema and data source generation working together.""" def test_flat_schema_and_data_alignment(self, flat_test_data): """Test that schema and data are aligned for flat structure.""" column_mapping = { "query": "${data.query}", "response": "${data.response}", "ground_truth": "${data.ground_truth}", } config = _generate_data_source_config(flat_test_data, column_mapping) data_source = _get_data_source(flat_test_data, column_mapping) schema_props = config["item_schema"]["properties"] data_item = data_source["source"]["content"][0][WRAPPER_KEY] # All schema properties should exist in data for prop_name in schema_props.keys(): assert prop_name in data_item def test_nested_schema_and_data_alignment(self, nested_test_data): """Test that schema and data are aligned for nested structure.""" column_mapping = { "query": "${data.item.query}", "rotation_days": "${data.item.context.company.policy.security.passwords.rotation_days}", "response": "${data.item.response}", } config = _generate_data_source_config(nested_test_data, column_mapping) data_source = _get_data_source(nested_test_data, column_mapping) # Both should handle nested structure consistently assert config["item_schema"]["type"] == "object" assert WRAPPER_KEY in data_source["source"]["content"][0] # Verify nested paths exist in data item = data_source["source"]["content"][0][WRAPPER_KEY] assert "query" in item assert "context" in item assert "company" in item["context"]