From 75a9327043476d2e8673bc5188976d5d4c9da128 Mon Sep 17 00:00:00 2001
From: Joseph Schull <joseph.schull@login-01.clusternet>
Date: Fri, 12 Jul 2024 11:17:40 -0700
Subject: [PATCH] Updated rename wells and testing

---
 iohub/cli/cli.py      | 18 ++++++++++++------
 iohub/ngff.py         | 10 ++++------
 tests/cli/test_cli.py | 32 +++++++++++++++++++++++++++++---
 3 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/iohub/cli/cli.py b/iohub/cli/cli.py
index 52f5dca5..993d3864 100644
--- a/iohub/cli/cli.py
+++ b/iohub/cli/cli.py
@@ -96,13 +96,15 @@ def convert(input, output, grid_layout, chunks):
 @click.option(
     "-i",
     "--input",
-    type=click.Path(exists=True, file_okay=True, dir_okay=False),
+    "zarrfile",
+    type=click.Path(exists=True, file_okay=True, dir_okay=True),
     required=True,
     help="Path to the input Zarr file.",
 )
 @click.option(
     "-c",
     "--csv",
+    "csvfile",
     type=click.File("r"),
     required=True,
     help="Path to the CSV file containing well names.",
@@ -116,6 +118,7 @@ def rename_wells_cli(csvfile, zarrfile):
 
     The CSV file should have two columns: old_well_path and new_well_path.
     """
+
     names = []
 
     csvreader = csv.reader(csvfile)
@@ -131,13 +134,16 @@ def rename_wells_cli(csvfile, zarrfile):
 
     modified = []
 
-    for old_well_path, new_well_path in names:
+    print(f"names: {names}")
+
+    for name in names:
+        old_well_path, new_well_path = name[0], name[1]
         for well in plate.metadata.wells:
-            if well.path == old_well_path and well not in modified:
+            if str(well.path) == str(old_well_path) and well not in modified:
+                print(f"Renaming {old_well_path} to {new_well_path}...")
                 try:
-                    plate.rename_well(
-                        well, old_well_path, new_well_path, modified, False
-                    )
+                    plate.rename_well(well, old_well_path, new_well_path)
                     modified.append(well)
+                    print(f"Well {old_well_path} renamed to {new_well_path}")
                 except ValueError as e:
                     click.echo(f"Error: {e}", err=True)
diff --git a/iohub/ngff.py b/iohub/ngff.py
index 460e00fb..23c143ee 100644
--- a/iohub/ngff.py
+++ b/iohub/ngff.py
@@ -1579,10 +1579,6 @@ def rename_well(
             Old name of well
         new_well_path : str
             New name of well
-        modified : dict, optional
-            Tracks modified wells
-        single well: bool, optional
-            True: renaming one well, False: renaming multiple wells
         """
         old_row, old_column = old_well_path.split("/")
         new_row, new_column = new_well_path.split("/")
@@ -1592,8 +1588,10 @@ def rename_well(
         if old_well_path not in well_paths:
             raise ValueError(f"Well '{old_well_path}' not found in plate.")
 
-        else:
-            well = self[old_well_path]
+        elif old_well_path in well_paths:
+            well = next(
+                w for w in self.metadata.wells if w.path == old_well_path
+            )
 
             if well.path == old_well_path:
                 well.path = new_well_path  # update well metadata
diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py
index c6e8ede0..19e980db 100644
--- a/tests/cli/test_cli.py
+++ b/tests/cli/test_cli.py
@@ -1,17 +1,17 @@
+import csv
 import re
 from unittest.mock import patch
 
-from click.testing import CliRunner
 import pytest
+from click.testing import CliRunner
 
 from iohub._version import __version__
 from iohub.cli.cli import cli
-
 from tests.conftest import (
+    hcs_ref,
     mm2gamma_ome_tiffs,
     ndtiff_v2_datasets,
     ndtiff_v3_labeled_positions,
-    hcs_ref,
 )
 
 
@@ -114,3 +114,29 @@ def test_rename_wells_help():
         result = runner.invoke(cli, cmd)
         assert result.exit_code == 0
         assert "containing well names" in result.output
+
+
+def test_rename_wells_basic():
+    runner = CliRunner()
+    test_zarr = "/hpc/mydata/joseph.schull/stitched_phase.zarr"
+    test_csv = "/hpc/mydata/joseph.schull/update_well_names.csv"
+    cmd = ["rename-wells", "-i", test_zarr, "-c", test_csv]
+    result = runner.invoke(cli, cmd)
+    print(result.output)
+    assert result.exit_code == 0
+
+
+def test_rename_wells_completion():
+    runner = CliRunner()
+    test_zarr = "/hpc/mydata/joseph.schull/stitched_phase.zarr"
+    test_csv = "/hpc/mydata/joseph.schull/update_well_names.csv"
+    cmd = ["rename-wells", "-i", test_zarr, "-c", test_csv]
+    result = runner.invoke(cli, cmd)
+    with open(test_csv, mode="r") as infile:
+        reader = csv.reader(infile)
+        names = list(reader)
+    for oldname, newname in names:
+        expected_message = f"Well {oldname} renamed to {newname}"
+        assert (
+            expected_message in result.output
+        ), f"Missing expected message: {expected_message}"