diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
new file mode 100644
index 0000000..f62a371
--- /dev/null
+++ b/.github/workflows/test.yaml
@@ -0,0 +1,46 @@
+name: test
+
+on:
+  push:
+    branches:
+      - master
+      - dev
+    tags:
+      - "v*"
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  test:
+    name: ${{ matrix.os }} ${{ matrix.python-version }}
+    runs-on: ${{ matrix.os }}
+    timeout-minutes: 90
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest]
+        python-version: ["3.11", "3.12"]
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Setup micromamba
+        uses: mamba-org/setup-micromamba@v2
+        with:
+          environment-file: environment.yaml
+          create-args: >-
+            python=${{ matrix.python-version }}
+            pytest
+
+      - name: Install slideio
+        shell: bash -el {0}
+        run: pip install slideio
+
+      - name: Install patho-sam
+        shell: bash -el {0}
+        run: pip install --no-deps -e .
+
+      - name: Run tests
+        shell: bash -el {0}
+        run: pytest test/ -v
diff --git a/patho_sam/automatic_segmentation.py b/patho_sam/automatic_segmentation.py
index f6316cc..5c77a77 100644
--- a/patho_sam/automatic_segmentation.py
+++ b/patho_sam/automatic_segmentation.py
@@ -111,7 +111,7 @@ def automatic_segmentation_wsi(
                 model_type=model_type,
                 checkpoint=checkpoint_path,
                 device=device,
-                amg=False,  # i.e. run AIS.
+                segmentation_mode="ais",  # i.e. run AIS.
                 is_tiled=isinstance(tile_shape, Tuple),  # i.e. run tiling-window based segmentation.
             )
 
@@ -125,7 +125,7 @@ def automatic_segmentation_wsi(
                 tile_shape=tile_shape,
                 halo=halo,
                 verbose=verbose,
-                output_mode=None,  # Skips some post-processing under `generate` method after automatic seg.
+                output_mode="instance_segmentation",  # Returns the raw instance segmentation array directly.
                 return_embeddings=True,  # Returns image embeddings, can be used in the task below, i.e. semantic seg.
                 batch_size=batch_size,
             )
diff --git a/test/__init__.py b/test/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/test/test_automatic_segmentation.py b/test/test_automatic_segmentation.py
new file mode 100644
index 0000000..fc81d48
--- /dev/null
+++ b/test/test_automatic_segmentation.py
@@ -0,0 +1,106 @@
+import os
+import unittest
+from shutil import rmtree
+
+import numpy as np
+from micro_sam.util import get_cache_directory
+from micro_sam.sample_data import fetch_wholeslide_histopathology_example_data
+
+
+DATA_CACHE = os.path.join(get_cache_directory(), "sample_data")
+
+
+class TestAutomaticSegmentationWSI(unittest.TestCase):
+    tmp_folder = "tmp-patho-sam-test"
+
+    @classmethod
+    def setUpClass(cls):
+        os.makedirs(cls.tmp_folder, exist_ok=True)
+        cls.wsi_path = fetch_wholeslide_histopathology_example_data(DATA_CACHE)
+        cls.roi = (0, 0, 512, 512)  # small crop for fast testing
+
+    @classmethod
+    def tearDownClass(cls):
+        rmtree(cls.tmp_folder, ignore_errors=True)
+
+    def test_instance_segmentation(self):
+        from patho_sam.automatic_segmentation import automatic_segmentation_wsi
+
+        output_path = os.path.join(self.tmp_folder, "instances.tif")
+        result = automatic_segmentation_wsi(
+            input_image=self.wsi_path,
+            model_type="vit_b_histopathology",
+            roi=self.roi,
+            output_path=output_path,
+            tile_shape=(384, 384),
+            halo=(64, 64),
+            output_choice="instances",
+            verbose=False,
+        )
+
+        self.assertEqual(result.ndim, 2)
+        self.assertEqual(result.shape, (self.roi[3], self.roi[2]))
+
+    def test_semantic_segmentation(self):
+        from patho_sam.automatic_segmentation import automatic_segmentation_wsi
+
+        output_path = os.path.join(self.tmp_folder, "semantic.tif")
+        result = automatic_segmentation_wsi(
+            input_image=self.wsi_path,
+            model_type="vit_b_histopathology",
+            roi=self.roi,
+            output_path=output_path,
+            tile_shape=(384, 384),
+            halo=(64, 64),
+            output_choice="semantic",
+            verbose=False,
+        )
+
+        self.assertEqual(result.ndim, 2)
+        self.assertEqual(result.shape, (self.roi[3], self.roi[2]))
+        # Class labels should be in range [0, 5] (6 PanNuke classes incl. background)
+        self.assertGreaterEqual(int(result.min()), 0)
+        self.assertLessEqual(int(result.max()), 5)
+
+    def test_all_segmentation(self):
+        from patho_sam.automatic_segmentation import automatic_segmentation_wsi
+
+        output_path = os.path.join(self.tmp_folder, "all.tif")
+        result = automatic_segmentation_wsi(
+            input_image=self.wsi_path,
+            model_type="vit_b_histopathology",
+            roi=self.roi,
+            output_path=output_path,
+            tile_shape=(384, 384),
+            halo=(64, 64),
+            output_choice="all",
+            verbose=False,
+        )
+
+        # "all" stacks instance + semantic → (2, H, W)
+        self.assertEqual(result.ndim, 3)
+        self.assertEqual(result.shape, (2, self.roi[3], self.roi[2]))
+
+    def test_instance_segmentation_idempotent(self):
+        """Running twice with the same output path should load from cache."""
+        from patho_sam.automatic_segmentation import automatic_segmentation_wsi
+
+        output_path = os.path.join(self.tmp_folder, "instances_cached.tif")
+        kwargs = dict(
+            input_image=self.wsi_path,
+            model_type="vit_b_histopathology",
+            roi=self.roi,
+            output_path=output_path,
+            tile_shape=(384, 384),
+            halo=(64, 64),
+            output_choice="instances",
+            verbose=False,
+        )
+
+        result1 = automatic_segmentation_wsi(**kwargs)
+        result2 = automatic_segmentation_wsi(**kwargs)
+        np.testing.assert_array_equal(result1, result2)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/test_cli.py b/test/test_cli.py
new file mode 100644
index 0000000..88e2008
--- /dev/null
+++ b/test/test_cli.py
@@ -0,0 +1,63 @@
+import os
+import subprocess
+import sys
+import unittest
+from shutil import rmtree
+
+from micro_sam.util import get_cache_directory
+from micro_sam.sample_data import fetch_wholeslide_histopathology_example_data
+
+
+DATA_CACHE = os.path.join(get_cache_directory(), "sample_data")
+
+
+def _run(args):
+    return subprocess.run([sys.executable, "-m"] + args, capture_output=True, text=True)
+
+
+class TestCLIRun(unittest.TestCase):
+    """End-to-end CLI run on a small WSI crop."""
+
+    tmp_folder = "tmp-patho-sam-cli-test"
+
+    @classmethod
+    def setUpClass(cls):
+        os.makedirs(cls.tmp_folder, exist_ok=True)
+        cls.wsi_path = fetch_wholeslide_histopathology_example_data(DATA_CACHE)
+
+    @classmethod
+    def tearDownClass(cls):
+        rmtree(cls.tmp_folder, ignore_errors=True)
+
+    def test_automatic_segmentation_instances(self):
+        output_path = os.path.join(self.tmp_folder, "cli_instances.tif")
+        result = _run([
+            "patho_sam.automatic_segmentation",
+            "-i", self.wsi_path,
+            "-o", output_path,
+            "--roi", "0", "0", "512", "512",
+            "-m", "vit_b_histopathology",
+            "--output_choice", "instances",
+        ])
+        self.assertEqual(result.returncode, 0, msg=result.stderr)
+        # The CLI writes <stem>_ROI_..._instances.tif
+        written = [f for f in os.listdir(self.tmp_folder) if "instances" in f]
+        self.assertTrue(len(written) > 0, "No instances output file found.")
+
+    def test_automatic_segmentation_semantic(self):
+        output_path = os.path.join(self.tmp_folder, "cli_semantic.tif")
+        result = _run([
+            "patho_sam.automatic_segmentation",
+            "-i", self.wsi_path,
+            "-o", output_path,
+            "--roi", "0", "0", "512", "512",
+            "-m", "vit_b_histopathology",
+            "--output_choice", "semantic",
+        ])
+        self.assertEqual(result.returncode, 0, msg=result.stderr)
+        written = [f for f in os.listdir(self.tmp_folder) if "semantic" in f]
+        self.assertTrue(len(written) > 0, "No semantic output file found.")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/test_io.py b/test/test_io.py
new file mode 100644
index 0000000..7e5dff1
--- /dev/null
+++ b/test/test_io.py
@@ -0,0 +1,53 @@
+import unittest
+
+
+class TestReadWSI(unittest.TestCase):
+
+    def test_file_not_found(self):
+        from patho_sam.io import read_wsi
+        with self.assertRaises(FileNotFoundError):
+            read_wsi("/nonexistent/path/image.svs")
+
+    @unittest.skipIf(
+        __import__("importlib").util.find_spec("slideio") is None, "slideio is not installed"
+    )
+    def test_read_wsi_example(self):
+        import os
+        from micro_sam.util import get_cache_directory
+        from micro_sam.sample_data import fetch_wholeslide_histopathology_example_data
+        from patho_sam.io import read_wsi
+
+        cache_dir = os.path.join(get_cache_directory(), "sample_data")
+        wsi_path = fetch_wholeslide_histopathology_example_data(cache_dir)
+
+        # Read a small ROI crop.
+        roi = (0, 0, 512, 512)
+        image = read_wsi(wsi_path, image_size=roi)
+
+        self.assertEqual(image.ndim, 3)  # H x W x C
+        self.assertEqual(image.shape[0], 512)
+        self.assertEqual(image.shape[1], 512)
+        self.assertEqual(image.shape[2], 3)  # RGB
+
+    @unittest.skipIf(
+        __import__("importlib").util.find_spec("slideio") is None, "slideio is not installed"
+    )
+    def test_read_wsi_with_scale(self):
+        import os
+        from micro_sam.util import get_cache_directory
+        from micro_sam.sample_data import fetch_wholeslide_histopathology_example_data
+        from patho_sam.io import read_wsi
+
+        cache_dir = os.path.join(get_cache_directory(), "sample_data")
+        wsi_path = fetch_wholeslide_histopathology_example_data(cache_dir)
+
+        # Read with a target scale.
+        image = read_wsi(wsi_path, scale=(256, 0))
+
+        self.assertEqual(image.ndim, 3)
+        # slideio preserves aspect ratio, so height may not be exactly 256
+        self.assertLessEqual(image.shape[0], 256)
+
+
+if __name__ == "__main__":
+    unittest.main()