From c40cf3514ee7c59e6d6a647b04288c6f2b051ac7 Mon Sep 17 00:00:00 2001
From: Nathan Fradet <56734983+Natooz@users.noreply.github.com>
Date: Thu, 6 Jul 2023 14:40:14 +0200
Subject: [PATCH] fixing data augmentation test to handle unique_track drum
 tokens

---
 tests/test_methods.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tests/test_methods.py b/tests/test_methods.py
index 6e6d4911..3d260359 100644
--- a/tests/test_methods.py
+++ b/tests/test_methods.py
@@ -214,10 +214,14 @@ def test_data_augmentation():
             ):
                 if is_drum:
                     continue
-                for original_token, aug_token in zip(original_track, aug_track):
+                for idx, (original_token, aug_token) in enumerate(zip(original_track, aug_track)):
                     if not tokenizer.is_multi_voc:
                         if original_token in pitch_tokens:
-                            assert aug_token == original_token + offsets[0]
+                            pitch_offset = offsets[0]
+                            # no offset for drum pitches
+                            if tokenizer.unique_track and idx > 0 and tokenizer[original_track[idx - 1]] == "Program_-1":
+                                pitch_offset = 0
+                            assert aug_token == original_token + pitch_offset
                         elif original_token in vel_tokens:
                             assert aug_token in [
                                 original_token + offsets[1],