"""actpred 模块单元测试 运行方式： cd actpred && python -m pytest test_actpred.py -v # 或 cd actpred && python test_actpred.py """ import os import tempfile import unittest from pathlib import Path import numpy as np from PIL import Image import torch from config import Config, DataConfig, ModelConfig, build_argparser, apply_overrides from dataset import _natural_key, ActionDataset, DataRecorder from model import ActionPredictor from predict import parse_action_history class TestNaturalKey(unittest.TestCase): """测试自然排序键函数""" def test_natural_order_no_padding(self): """无前导零的文件名应按数字序排列""" files = [Path(f"{i}.jpg") for i in [3, 1, 10, 2, 20]] sorted_files = sorted(files, key=_natural_key) names = [f.name for f in sorted_files] self.assertEqual(names, ["1.jpg", "2.jpg", "3.jpg", "10.jpg", "20.jpg"]) def test_natural_order_with_padding(self): """有前导零的文件名排序也应正确""" files = [Path(f"{i:04d}.jpg") for i in [3, 1, 10, 2, 20]] sorted_files = sorted(files, key=_natural_key) names = [f.name for f in sorted_files] self.assertEqual(names, ["0001.jpg", "0002.jpg", "0003.jpg", "0010.jpg", "0020.jpg"]) class TestParseActionHistory(unittest.TestCase): """测试动作历史解析函数""" def test_single_step(self): history = parse_action_history("1,2,3,4", action_dim=4) self.assertEqual(history.shape, (1, 4)) np.testing.assert_array_equal(history[0], [1, 2, 3, 4]) def test_multi_step(self): history = parse_action_history("1,0,0,0:0,1,0,0:0,0,1,0", action_dim=4) self.assertEqual(history.shape, (3, 4)) def test_wrong_dim_raises(self): with self.assertRaises(AssertionError): parse_action_history("1,2,3", action_dim=4) class TestActionPredictor(unittest.TestCase): """测试模型前向传播""" def test_forward_simple_mlp(self): data_cfg = DataConfig(image_size=64, action_dim=8, history_len=5) model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="mlp") model = ActionPredictor(model_cfg, data_cfg) model.eval() image = torch.randn(2, 3, 64, 64) history = torch.randn(2, 5, 8) with torch.no_grad(): out = model(image, history) self.assertEqual(out.shape, (2, 8)) def test_forward_simple_gru(self): data_cfg = DataConfig(image_size=64, action_dim=8, history_len=5) model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="gru") model = ActionPredictor(model_cfg, data_cfg) model.eval() image = torch.randn(2, 3, 64, 64) history = torch.randn(2, 5, 8) with torch.no_grad(): out = model(image, history) self.assertEqual(out.shape, (2, 8)) def test_forward_dynamic_batch(self): data_cfg = DataConfig(image_size=32, action_dim=4, history_len=3) model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="mlp") model = ActionPredictor(model_cfg, data_cfg) model.eval() for batch_size in [1, 4, 8]: image = torch.randn(batch_size, 3, 32, 32) history = torch.randn(batch_size, 3, 4) with torch.no_grad(): out = model(image, history) self.assertEqual(out.shape, (batch_size, 4)) class TestDataRecorder(unittest.TestCase): """测试数据录制器""" def test_record_and_flush(self): with tempfile.TemporaryDirectory() as tmpdir: recorder = DataRecorder(tmpdir, action_dim=4) for _ in range(10): img = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8) action = np.random.randn(4).astype(np.float32) recorder.record(img, action) episode_path = recorder.flush() self.assertTrue(os.path.exists(os.path.join(episode_path, "actions.npy"))) images_dir = os.path.join(episode_path, "images") self.assertTrue(os.path.isdir(images_dir)) self.assertEqual(len(os.listdir(images_dir)), 10) actions = np.load(os.path.join(episode_path, "actions.npy")) self.assertEqual(actions.shape, (10, 4)) def test_multiple_episodes_increment_index(self): with tempfile.TemporaryDirectory() as tmpdir: r1 = DataRecorder(tmpdir, action_dim=4) r1.flush() r2 = DataRecorder(tmpdir, action_dim=4) r2.flush() self.assertNotEqual(r1.episode_dir, r2.episode_dir) self.assertIn("episode_0001", str(r2.episode_dir)) class TestActionDataset(unittest.TestCase): """测试数据集加载""" @staticmethod def _create_dummy_episode(base_dir, n_frames=10, action_dim=4, image_size=32): ep_dir = base_dir / "episode_0001" images_dir = ep_dir / "images" images_dir.mkdir(parents=True, exist_ok=True) actions = [] for i in range(n_frames): img = Image.fromarray( np.random.randint(0, 255, (image_size, image_size, 3), dtype=np.uint8)) img.save(images_dir / f"{i:04d}.jpg") actions.append(np.random.randn(action_dim).astype(np.float32)) np.save(ep_dir / "actions.npy", np.array(actions)) def test_dataset_loading(self): with tempfile.TemporaryDirectory() as tmpdir: self._create_dummy_episode(Path(tmpdir), n_frames=10, action_dim=4) ds = ActionDataset(tmpdir, history_len=3, action_dim=4, image_size=32) self.assertEqual(len(ds), 7) # 10 - 3 = 7 samples sample = ds[0] self.assertIn("image", sample) self.assertIn("action_history", sample) self.assertIn("target", sample) self.assertEqual(sample["image"].shape, (3, 32, 32)) self.assertEqual(sample["action_history"].shape, (3, 4)) self.assertEqual(sample["target"].shape, (4,)) def test_episode_count(self): with tempfile.TemporaryDirectory() as tmpdir: self._create_dummy_episode(Path(tmpdir)) count = ActionDataset.get_episode_count(tmpdir) self.assertEqual(count, 1) def test_empty_stats_raises(self): """空数据集计算归一化统计量应报错""" with tempfile.TemporaryDirectory() as tmpdir: with self.assertRaises(ValueError): ActionDataset(tmpdir, history_len=5, action_dim=4, normalize_actions=True) def test_invalid_stats_path_raises(self): """指定的 stats 路径不存在应报错""" with tempfile.TemporaryDirectory() as tmpdir: with self.assertRaises(FileNotFoundError): ActionDataset(tmpdir, history_len=3, action_dim=4, normalize_actions=True, action_stats_path="/nonexistent/path.npy") class TestConfig(unittest.TestCase): """测试配置和 CLI 参数""" def test_default_config(self): cfg = Config() self.assertEqual(cfg.data.image_size, 64) self.assertEqual(cfg.data.action_dim, 8) self.assertEqual(cfg.data.history_len, 5) self.assertEqual(cfg.model.cnn_backbone, "simple") self.assertEqual(cfg.model.history_encoder, "mlp") self.assertEqual(cfg.train.epochs, 100) def test_apply_overrides(self): parser = build_argparser("test") args = parser.parse_args(["--epochs", "50", "--lr", "0.01", "--backbone", "mobilenet"]) cfg = apply_overrides(Config(), args) self.assertEqual(cfg.train.epochs, 50) self.assertEqual(cfg.train.lr, 0.01) self.assertEqual(cfg.model.cnn_backbone, "mobilenet") def test_normalize_actions_flag(self): """--normalize-actions 和 --no-normalize-actions 都应可用""" parser = build_argparser("test") args = parser.parse_args(["--normalize-actions"]) self.assertTrue(args.normalize_actions) parser = build_argparser("test") args = parser.parse_args(["--no-normalize-actions"]) self.assertFalse(args.normalize_actions) parser = build_argparser("test") args = parser.parse_args([]) self.assertIsNone(args.normalize_actions) if __name__ == "__main__": unittest.main(verbosity=2)