tests/test_segmentation.py

import unittest
import cv2
import numpy as np
from app.models.segmenter import PuzzlePieceSegmenter

class TestSegmentation(unittest.TestCase):
    def setUp(self):
        self.segmenter = PuzzlePieceSegmenter()
        self.test_image = np.zeros((1000, 1000, 3), dtype=np.uint8)  # Dummy image

    def test_segment_pieces(self):
        pieces = self.segmenter.segment_pieces(self.test_image)
        self.assertIsInstance(pieces, list)
        for piece in pieces:
            self.assertIn('id', piece)
            self.assertIn('image', piece)
            self.assertIn('features', piece)

if __name__ == '__main__':
    unittest.main()