From d80b4a11d7af280c4a97b12d6feb56c35b85d9c8 Mon Sep 17 00:00:00 2001
From: Matt Stone <matt@fulcrumgenomics.com>
Date: Tue, 17 Oct 2023 11:11:36 -0400
Subject: [PATCH] fix: tag value types, two char tags, and tests

---
 fgpyo/sam/__init__.py                     |  6 ++++-
 fgpyo/sam/tests/test_template_iterator.py | 30 +++++++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/fgpyo/sam/__init__.py b/fgpyo/sam/__init__.py
index d793d052..1f8a1e64 100644
--- a/fgpyo/sam/__init__.py
+++ b/fgpyo/sam/__init__.py
@@ -846,15 +846,19 @@ def all_recs(self) -> Iterator[AlignedSegment]:
     def set_tag(
         self,
         tag: str,
-        value: Union[str, int, float],
+        value: Union[str, int, float, None],
     ) -> None:
         """Add a tag to all records associated with the template.
 
+        Setting a tag to `None` will remove the tag.
+
         Args:
             tag: The name of the tag.
             value: The value of the tag.
         """
 
+        assert len(tag) == 2, "Tags must be 2 characters."
+
         for rec in self.all_recs():
             rec.set_tag(tag, value)
 
diff --git a/fgpyo/sam/tests/test_template_iterator.py b/fgpyo/sam/tests/test_template_iterator.py
index cdefc35a..5c060355 100644
--- a/fgpyo/sam/tests/test_template_iterator.py
+++ b/fgpyo/sam/tests/test_template_iterator.py
@@ -1,3 +1,5 @@
+import pytest
+
 from fgpyo.sam import Template
 from fgpyo.sam.builder import SamBuilder
 
@@ -64,3 +66,31 @@ def test_to_templates() -> None:
     assert len(template2.r2_secondaries) == 0
     assert len(list(template2.primary_recs())) == 1
     assert len(list(template2.all_recs())) == 2
+
+
+def test_set_tag() -> None:
+    builder = SamBuilder()
+    template = Template.build(builder.add_pair(chrom="chr1", start1=100, start2=200))
+
+    TAG = "XF"
+    VALUE = "value"
+
+    for read in template.all_recs():
+        with pytest.raises(KeyError):
+            read.get_tag(TAG)
+
+    # test setting
+    template.set_tag(TAG, VALUE)
+    assert template.r1.get_tag(TAG) == VALUE
+    assert template.r2.get_tag(TAG) == VALUE
+
+    # test removal
+    template.set_tag(TAG, None)
+    for read in template.all_recs():
+        with pytest.raises(KeyError):
+            read.get_tag(TAG)
+
+    # test tags that aren't two characters
+    for bad_tag in ["", "A", "ABC", "ABCD"]:
+        with pytest.raises(AssertionError, match="Tags must be 2 characters."):
+            template.set_tag(bad_tag, VALUE)