import pytest import numpy as np import math from datetime import datetime from unittest.mock import Mock from f06_image_rotation_manager import ( ImageRotationManager, IImageMatcher, AlignmentResult, ChunkAlignmentResult, RotationResult ) import h07_image_rotation_utils as H07 @pytest.fixture def rm(): return ImageRotationManager() @pytest.fixture def dummy_image(): # Create a simple 100x100 asymmetrical image to easily test rotations img = np.zeros((100, 100, 3), dtype=np.uint8) img[10:30, 10:30] = [255, 255, 255] # White square in top-left return img @pytest.fixture def mock_matcher(): matcher = Mock(spec=IImageMatcher) # Default behavior: fail matcher.align_to_satellite.return_value = AlignmentResult(matched=False, confidence=0.0, homography=None, inlier_count=0) matcher.align_chunk_to_satellite.return_value = ChunkAlignmentResult(matched=False, confidence=0.0, homography=None, inlier_count=0) return matcher class TestImageRotationManager: # --- Feature 06.01: Image Rotation Core --- def test_rotate_image_360_identity(self, rm, dummy_image): rotated = rm.rotate_image_360(dummy_image, 0.0) assert np.array_equal(rotated, dummy_image) rotated_360 = rm.rotate_image_360(dummy_image, 360.0) assert np.array_equal(rotated_360, dummy_image) def test_rotate_image_360_angles(self, rm, dummy_image): # Test 90 degrees clockwise (white square moves from top-left to top-right) rotated_90 = rm.rotate_image_360(dummy_image, 90.0) assert rotated_90[10:30, 70:90].sum() > 0 # White region shifted # Dimensions preserved assert rotated_90.shape == dummy_image.shape # Rotate 180 degrees rotated_180 = rm.rotate_image_360(dummy_image, 180.0) assert rotated_180[70:90, 70:90].sum() > 0 # White region shifted to bottom-right # Rotate 270 degrees rotated_270 = rm.rotate_image_360(dummy_image, 270.0) assert rotated_270[70:90, 10:30].sum() > 0 # White region shifted to bottom-left # Rotate 45 degrees rotated_45 = rm.rotate_image_360(dummy_image, 45.0) assert rotated_45.shape == dummy_image.shape # Negative normalization rotated_neg = rm.rotate_image_360(dummy_image, -90.0) assert np.array_equal(rotated_neg, rotated_270) # Large angle normalization (450 = 90) rotated_450 = rm.rotate_image_360(dummy_image, 450.0) assert np.array_equal(rotated_450, rotated_90) def test_rotate_chunk_360(self, rm, dummy_image): # Empty chunk assert rm.rotate_chunk_360([], 90.0) == [] # Single image chunk single_chunk = [dummy_image.copy()] rotated_single = rm.rotate_chunk_360(single_chunk, 90.0) assert len(rotated_single) == 1 assert np.array_equal(rotated_single[0], rm.rotate_image_360(dummy_image, 90.0)) # Multiple images chunk = [dummy_image.copy(), dummy_image.copy()] rotated_chunk = rm.rotate_chunk_360(chunk, 180.0) assert len(rotated_chunk) == 2 # White square should be at bottom-right for both assert rotated_chunk[0][70:90, 70:90].sum() > 0 assert rotated_chunk[1][70:90, 70:90].sum() > 0 # Independence - original unchanged assert chunk[0][10:30, 10:30].sum() > 0 # --- Feature 06.02: Heading Management --- def test_internal_heading_helpers(self, rm): # _normalize_angle assert rm._normalize_angle(370.0) == 10.0 assert rm._normalize_angle(-30.0) == 330.0 # _calculate_angle_delta assert rm._calculate_angle_delta(10.0, 20.0) == 10.0 assert rm._calculate_angle_delta(350.0, 10.0) == 20.0 # Wraparound assert rm._calculate_angle_delta(10.0, 350.0) == 20.0 # Commutative assert rm._calculate_angle_delta(0.0, 180.0) == 180.0 # _get_flight_state (empty initially) assert rm._get_flight_state("flight_none") is None def test_heading_lifecycle(self, rm): ts = datetime.utcnow() # 1. New flight assert rm.get_current_heading("flight_1") is None assert rm.requires_rotation_sweep("flight_1") is True # 2. Update heading assert rm.update_heading("flight_1", 1, 45.0, ts) is True assert rm.get_current_heading("flight_1") == 45.0 assert rm.requires_rotation_sweep("flight_1") is False # 3. Angle normalization rm.update_heading("flight_1", 2, 370.0, ts) assert rm.get_current_heading("flight_1") == 10.0 rm.update_heading("flight_1", 3, -30.0, ts) assert rm.get_current_heading("flight_1") == 330.0 # 4. History limit for i in range(15): rm.update_heading("flight_1", i, float(i), ts) state = rm._get_flight_state("flight_1") assert len(state.heading_history) == rm.config.history_size assert state.heading_history[-1] == 14.0 def test_detect_sharp_turn(self, rm): ts = datetime.utcnow() # No heading yet -> cannot detect turn assert rm.detect_sharp_turn("flight_turn", 60.0) is False rm.update_heading("flight_turn", 1, 60.0, ts) assert rm.detect_sharp_turn("flight_turn", 75.0) is False # Small 15 deg assert rm.detect_sharp_turn("flight_turn", 120.0) is True # Sharp 60 deg assert rm.detect_sharp_turn("flight_turn", 105.0) is False # Exactly 45 threshold assert rm.detect_sharp_turn("flight_turn", 106.0) is True # Exactly 46 (> 45) # Wraparound testing rm.update_heading("flight_wrap", 1, 350.0, ts) assert rm.detect_sharp_turn("flight_wrap", 20.0) is False # 30 deg distance assert rm.detect_sharp_turn("flight_wrap", 60.0) is True # 70 deg distance def test_requires_rotation_sweep_flag(self, rm): ts = datetime.utcnow() rm.update_heading("flight_flag", 1, 90.0, ts) assert rm.requires_rotation_sweep("flight_flag") is False rm._set_sweep_required("flight_flag", True) assert rm.requires_rotation_sweep("flight_flag") is True def test_integration_heading_lifecycle(self, rm): # Test 1: Heading Lifecycle ts = datetime.utcnow() assert rm.get_current_heading("flight_int_1") is None assert rm.requires_rotation_sweep("flight_int_1") is True rm.update_heading("flight_int_1", 1, 45.0, ts) assert rm.get_current_heading("flight_int_1") == 45.0 assert rm.requires_rotation_sweep("flight_int_1") is False def test_integration_sharp_turn_flow(self, rm): # Test 2: Sharp Turn Flow ts = datetime.utcnow() rm.update_heading("flight_int_2", 1, 90.0, ts) assert rm.detect_sharp_turn("flight_int_2", 100.0) is False assert rm.detect_sharp_turn("flight_int_2", 180.0) is True rm._set_sweep_required("flight_int_2", True) assert rm.requires_rotation_sweep("flight_int_2") is True # --- Feature 06.03: Rotation Sweep Orchestration --- def _build_mock_homography(self, angle_deg: float) -> np.ndarray: rad = math.radians(angle_deg) return np.array([ [math.cos(rad), -math.sin(rad), 0], [math.sin(rad), math.cos(rad), 0], [0, 0, 1] ]) def test_internal_orchestration_helpers(self, rm): # _get_rotation_steps steps = rm._get_rotation_steps() assert len(steps) == 12 assert 0.0 in steps and 330.0 in steps # _extract_rotation_from_homography & _combine_angles H = self._build_mock_homography(10.0) delta = rm._extract_rotation_from_homography(H) assert abs(delta - 10.0) < 0.01 combined = rm._combine_angles(355.0, 10.0) assert combined == 5.0 # _select_best_result res1 = AlignmentResult(matched=True, confidence=0.8, homography=None, inlier_count=10) res2 = AlignmentResult(matched=True, confidence=0.95, homography=None, inlier_count=20) res3 = AlignmentResult(matched=False, confidence=0.99, homography=None, inlier_count=0) # Unmatched results = [(0.0, res1), (30.0, res2), (60.0, res3)] best = rm._select_best_result(results) assert best is not None assert best[0] == 30.0 # highest confidence that is also matched # Below threshold rm.config.confidence_threshold = 0.99 assert rm._select_best_result([(0.0, res1), (30.0, res2)]) is None def test_calculate_precise_angle(self, rm): # Small positive rotation H_pos = self._build_mock_homography(2.5) assert abs(rm.calculate_precise_angle(H_pos, 60.0) - 62.5) < 0.01 # Negative rotation H_neg = self._build_mock_homography(-3.0) assert abs(rm.calculate_precise_angle(H_neg, 30.0) - 27.0) < 0.01 # Wraparound large rotation H_wrap = self._build_mock_homography(20.0) assert abs(rm.calculate_precise_angle(H_wrap, 350.0) - 10.0) < 0.01 # Invalid homography fallback assert rm.calculate_precise_angle(None, 60.0) == 60.0 def test_try_rotation_steps_match_found(self, rm, dummy_image, mock_matcher): ts = datetime.utcnow() def mock_align(rotated_img, sat, bounds): # Simulate match ONLY when called with the 60 degree rotation # Since we iterate 0, 30, 60... this should happen on the 3rd step if mock_matcher.align_to_satellite.call_count == 3: H = self._build_mock_homography(2.3) return AlignmentResult(matched=True, confidence=0.85, homography=H, inlier_count=50) return AlignmentResult(matched=False, confidence=0.0, homography=None, inlier_count=0) mock_matcher.align_to_satellite.side_effect = mock_align res = rm.try_rotation_steps("flight_sweep", 1, dummy_image, dummy_image, None, ts, mock_matcher) assert res is not None assert res.matched is True assert res.initial_angle == 60.0 assert abs(res.precise_angle - 62.3) < 0.01 # State tracking updated correctly assert rm.get_current_heading("flight_sweep") == res.precise_angle def test_try_rotation_steps_no_match(self, rm, dummy_image, mock_matcher): ts = datetime.utcnow() # Matcher returns False for all (default mock behavior) res = rm.try_rotation_steps("flight_fail", 1, dummy_image, dummy_image, None, ts, mock_matcher) assert res is None assert mock_matcher.align_to_satellite.call_count == 12 # Exhausted all 12 angles assert rm.get_current_heading("flight_fail") is None # Unchanged def test_try_chunk_rotation_steps(self, rm, dummy_image, mock_matcher): chunk = [dummy_image, dummy_image] def mock_align_chunk(rotated_chunk, sat, bounds): # Simulate match at 90 degrees if mock_matcher.align_chunk_to_satellite.call_count == 4: return ChunkAlignmentResult(matched=True, confidence=0.9, homography=self._build_mock_homography(0), inlier_count=100) return ChunkAlignmentResult(matched=False, confidence=0.0, homography=None, inlier_count=0) mock_matcher.align_chunk_to_satellite.side_effect = mock_align_chunk res = rm.try_chunk_rotation_steps(chunk, dummy_image, None, mock_matcher) assert res is not None assert res.initial_angle == 90.0 assert res.precise_angle == 90.0 # Ensure mock was passed chunks of length 2 during its attempts args, _ = mock_matcher.align_chunk_to_satellite.call_args assert len(args[0]) == 2