""" Unit tests for TrainingDataExporter Tests the training data export functionality for AtomizerField. """ import pytest import json import tempfile import shutil from pathlib import Path from optimization_engine.training_data_exporter import TrainingDataExporter, create_exporter_from_config class TestTrainingDataExporter: """Test suite for TrainingDataExporter class.""" @pytest.fixture def temp_dir(self): """Create a temporary directory for tests.""" temp_path = Path(tempfile.mkdtemp()) yield temp_path # Cleanup after test if temp_path.exists(): shutil.rmtree(temp_path) @pytest.fixture def sample_exporter(self, temp_dir): """Create a sample exporter for testing.""" return TrainingDataExporter( export_dir=temp_dir / "test_export", study_name="test_study", design_variable_names=["thickness", "width", "length"], objective_names=["max_stress", "mass"], constraint_names=["stress_constraint"], metadata={"test_key": "test_value"} ) @pytest.fixture def sample_simulation_files(self, temp_dir): """Create sample .dat and .op2 files for testing.""" sim_dir = temp_dir / "sim_files" sim_dir.mkdir(parents=True, exist_ok=True) dat_file = sim_dir / "test_model.dat" op2_file = sim_dir / "test_model.op2" # Create dummy files with some content dat_file.write_text("$ Nastran input deck\nGRID,1,0,0.0,0.0,0.0\n") op2_file.write_bytes(b"DUMMY OP2 BINARY DATA" * 100) return { "dat_file": dat_file, "op2_file": op2_file } def test_exporter_initialization(self, temp_dir): """Test that exporter initializes correctly.""" export_dir = temp_dir / "test_export" exporter = TrainingDataExporter( export_dir=export_dir, study_name="beam_study", design_variable_names=["thickness", "width"], objective_names=["stress", "mass"] ) assert exporter.export_dir == export_dir assert exporter.study_name == "beam_study" assert exporter.design_variable_names == ["thickness", "width"] assert exporter.objective_names == ["stress", "mass"] assert exporter.trial_count == 0 assert len(exporter.exported_trials) == 0 assert export_dir.exists() def test_readme_creation(self, sample_exporter): """Test that README.md is created.""" readme_path = sample_exporter.export_dir / "README.md" assert readme_path.exists() content = readme_path.read_text() assert "test_study" in content assert "AtomizerField Training Data" in content assert "thickness" in content assert "max_stress" in content def test_export_trial_success(self, sample_exporter, sample_simulation_files): """Test successful trial export.""" design_vars = {"thickness": 3.5, "width": 50.0, "length": 200.0} results = { "objectives": {"max_stress": 245.3, "mass": 1.25}, "constraints": {"stress_constraint": -54.7}, "max_displacement": 1.23 } success = sample_exporter.export_trial( trial_number=1, design_variables=design_vars, results=results, simulation_files=sample_simulation_files ) assert success assert sample_exporter.trial_count == 1 assert len(sample_exporter.exported_trials) == 1 # Check directory structure trial_dir = sample_exporter.export_dir / "trial_0001" assert trial_dir.exists() assert (trial_dir / "input").exists() assert (trial_dir / "output").exists() assert (trial_dir / "input" / "model.bdf").exists() assert (trial_dir / "output" / "model.op2").exists() assert (trial_dir / "metadata.json").exists() def test_export_trial_metadata(self, sample_exporter, sample_simulation_files): """Test that metadata.json is created correctly.""" design_vars = {"thickness": 3.5, "width": 50.0, "length": 200.0} results = { "objectives": {"max_stress": 245.3, "mass": 1.25}, "constraints": {"stress_constraint": -54.7}, "max_displacement": 1.23 } sample_exporter.export_trial( trial_number=42, design_variables=design_vars, results=results, simulation_files=sample_simulation_files ) metadata_path = sample_exporter.export_dir / "trial_0042" / "metadata.json" assert metadata_path.exists() with open(metadata_path, 'r') as f: metadata = json.load(f) assert metadata["trial_number"] == 42 assert metadata["atomizer_study"] == "test_study" assert metadata["design_parameters"] == design_vars assert metadata["results"]["objectives"] == {"max_stress": 245.3, "mass": 1.25} assert metadata["results"]["constraints"] == {"stress_constraint": -54.7} assert metadata["results"]["max_displacement"] == 1.23 assert "timestamp" in metadata def test_export_trial_missing_dat_file(self, sample_exporter, temp_dir): """Test export fails gracefully when .dat file is missing.""" design_vars = {"thickness": 3.5} results = {"objectives": {"max_stress": 245.3}} # Create only .op2 file sim_dir = temp_dir / "sim_files" sim_dir.mkdir(parents=True, exist_ok=True) op2_file = sim_dir / "test_model.op2" op2_file.write_bytes(b"DUMMY OP2 DATA") simulation_files = { "dat_file": sim_dir / "nonexistent.dat", # Does not exist "op2_file": op2_file } success = sample_exporter.export_trial( trial_number=1, design_variables=design_vars, results=results, simulation_files=simulation_files ) assert not success assert sample_exporter.trial_count == 0 def test_export_trial_missing_op2_file(self, sample_exporter, temp_dir): """Test export fails gracefully when .op2 file is missing.""" design_vars = {"thickness": 3.5} results = {"objectives": {"max_stress": 245.3}} # Create only .dat file sim_dir = temp_dir / "sim_files" sim_dir.mkdir(parents=True, exist_ok=True) dat_file = sim_dir / "test_model.dat" dat_file.write_text("$ Nastran input") simulation_files = { "dat_file": dat_file, "op2_file": sim_dir / "nonexistent.op2" # Does not exist } success = sample_exporter.export_trial( trial_number=1, design_variables=design_vars, results=results, simulation_files=simulation_files ) assert not success assert sample_exporter.trial_count == 0 def test_export_multiple_trials(self, sample_exporter, sample_simulation_files): """Test exporting multiple trials.""" for i in range(1, 6): design_vars = {"thickness": 2.0 + i * 0.5, "width": 40.0 + i * 5, "length": 180.0 + i * 10} results = { "objectives": {"max_stress": 200.0 + i * 10, "mass": 1.0 + i * 0.1} } success = sample_exporter.export_trial( trial_number=i, design_variables=design_vars, results=results, simulation_files=sample_simulation_files ) assert success assert sample_exporter.trial_count == 5 assert len(sample_exporter.exported_trials) == 5 # Verify all trial directories exist for i in range(1, 6): trial_dir = sample_exporter.export_dir / f"trial_{i:04d}" assert trial_dir.exists() assert (trial_dir / "metadata.json").exists() def test_finalize(self, sample_exporter, sample_simulation_files): """Test finalize creates study_summary.json.""" # Export a few trials for i in range(1, 4): design_vars = {"thickness": 2.0 + i * 0.5, "width": 50.0, "length": 200.0} results = {"objectives": {"max_stress": 200.0 + i * 10, "mass": 1.2}} sample_exporter.export_trial(i, design_vars, results, sample_simulation_files) # Finalize sample_exporter.finalize() # Check study_summary.json summary_path = sample_exporter.export_dir / "study_summary.json" assert summary_path.exists() with open(summary_path, 'r') as f: summary = json.load(f) assert summary["study_name"] == "test_study" assert summary["total_trials"] == 3 assert summary["design_variables"] == ["thickness", "width", "length"] assert summary["objectives"] == ["max_stress", "mass"] assert summary["constraints"] == ["stress_constraint"] assert "export_timestamp" in summary assert summary["metadata"] == {"test_key": "test_value"} def test_create_exporter_from_config_disabled(self): """Test that None is returned when export is disabled.""" config = { "training_data_export": { "enabled": False, "export_dir": "/some/path" } } exporter = create_exporter_from_config(config) assert exporter is None def test_create_exporter_from_config_missing_export_dir(self): """Test that None is returned when export_dir is missing.""" config = { "training_data_export": { "enabled": True } } exporter = create_exporter_from_config(config) assert exporter is None def test_create_exporter_from_config_success(self, temp_dir): """Test successful exporter creation from config.""" config = { "study_name": "beam_optimization", "training_data_export": { "enabled": True, "export_dir": str(temp_dir / "export") }, "design_variables": [ {"name": "thickness"}, {"parameter": "width"}, {"other_key": "value"} # Should use "var_2" ], "objectives": [ {"name": "max_stress"}, {"name": "mass"} ], "constraints": [ {"name": "stress_limit"} ], "version": "1.0", "optimization": { "algorithm": "NSGA-II", "n_trials": 100 } } exporter = create_exporter_from_config(config) assert exporter is not None assert exporter.study_name == "beam_optimization" assert exporter.design_variable_names == ["thickness", "width", "var_2"] assert exporter.objective_names == ["max_stress", "mass"] assert exporter.constraint_names == ["stress_limit"] assert exporter.study_metadata["atomizer_version"] == "1.0" assert exporter.study_metadata["optimization_algorithm"] == "NSGA-II" assert exporter.study_metadata["n_trials"] == 100 def test_create_exporter_from_config_defaults(self, temp_dir): """Test exporter creation with minimal config.""" config = { "training_data_export": { "enabled": True, "export_dir": str(temp_dir / "export") } } exporter = create_exporter_from_config(config) assert exporter is not None assert exporter.study_name == "unnamed_study" assert exporter.design_variable_names == [] assert exporter.objective_names == [] assert exporter.constraint_names == [] def test_file_content_preservation(self, sample_exporter, temp_dir): """Test that file contents are preserved during copy.""" # Create files with specific content sim_dir = temp_dir / "sim_files" sim_dir.mkdir(parents=True, exist_ok=True) dat_content = "$ Nastran Input Deck\nGRID,1,0,0.0,0.0,0.0\nGRID,2,0,1.0,0.0,0.0\n" op2_content = b"\x00\x01\x02\x03NASTRAN_BINARY_DATA\xFF\xFE\xFD\xFC" * 50 dat_file = sim_dir / "model.dat" op2_file = sim_dir / "model.op2" dat_file.write_text(dat_content) op2_file.write_bytes(op2_content) simulation_files = {"dat_file": dat_file, "op2_file": op2_file} # Export trial sample_exporter.export_trial( trial_number=1, design_variables={"thickness": 3.0}, results={"objectives": {"stress": 200.0}}, simulation_files=simulation_files ) # Verify copied files have same content exported_dat = sample_exporter.export_dir / "trial_0001" / "input" / "model.bdf" exported_op2 = sample_exporter.export_dir / "trial_0001" / "output" / "model.op2" assert exported_dat.read_text() == dat_content assert exported_op2.read_bytes() == op2_content def test_trial_numbering_format(self, sample_exporter, sample_simulation_files): """Test that trial directories use correct numbering format.""" # Test various trial numbers trial_numbers = [1, 9, 10, 99, 100, 999, 1000] for trial_num in trial_numbers: sample_exporter.export_trial( trial_number=trial_num, design_variables={"thickness": 3.0}, results={"objectives": {"stress": 200.0}}, simulation_files=sample_simulation_files ) expected_dir = sample_exporter.export_dir / f"trial_{trial_num:04d}" assert expected_dir.exists() # Verify specific formatting assert (sample_exporter.export_dir / "trial_0001").exists() assert (sample_exporter.export_dir / "trial_0009").exists() assert (sample_exporter.export_dir / "trial_0010").exists() assert (sample_exporter.export_dir / "trial_0099").exists() assert (sample_exporter.export_dir / "trial_0100").exists() assert (sample_exporter.export_dir / "trial_0999").exists() assert (sample_exporter.export_dir / "trial_1000").exists() if __name__ == "__main__": pytest.main([__file__, "-v"])