From eab3517d5d8db6690177b0e7b360a6c14cd2dd70 Mon Sep 17 00:00:00 2001
From: Valerij Talagayev <82884038+talagayev@users.noreply.github.com>
Date: Mon, 16 Dec 2024 13:55:17 +0100
Subject: [PATCH] Create test_barcode_generation.py

added first pytests
---
 .../test_barcode_generation.py                | 41 +++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 openmmdl/tests/openmmdl_analysis/test_barcode_generation.py

diff --git a/openmmdl/tests/openmmdl_analysis/test_barcode_generation.py b/openmmdl/tests/openmmdl_analysis/test_barcode_generation.py
new file mode 100644
index 0000000..ef75fa8
--- /dev/null
+++ b/openmmdl/tests/openmmdl_analysis/test_barcode_generation.py
@@ -0,0 +1,41 @@
+import numpy as np
+import pandas as pd
+import re
+import os
+import matplotlib.pyplot as plt
+import pytest
+from openmmdl.openmmdl_analysis.barcode_generation import BarcodeGenerator, BarcodePlotter
+
+# Barcode generation tests
+@pytest.fixture
+def sample_dataframe_barcode_generation():
+    data = {
+        'FRAME': [1, 1, 2, 2, 3],
+        'Interaction1': [1, 0, 1, 0, 0],
+        'Interaction2': [0, 0, 0, 1, 1],
+        'WATER_IDX': [101, 102, 103, 104, 105],
+    }
+    return pd.DataFrame(data)
+
+def test_barcodegeneration(sample_dataframe_barcode_generation):
+    interaction = 'Interaction1'
+    barcode_generator = BarcodeGenerator(sample_dataframe_barcode_generation)
+
+    barcode = barcode_generator.generate_barcode(interaction)
+    
+    assert isinstance(barcode, np.ndarray)
+    
+    expected_barcode = np.array([1, 1, 0])
+    assert np.array_equal(barcode, expected_barcode)
+    
+def test_waterids_barcode_generator(sample_dataframe_barcode_generation):
+    interaction = 'Interaction2'
+    barcode_generator = BarcodeGenerator(sample_dataframe_barcode_generation)
+    waterid_barcode = barcode_generator.generate_waterids_barcode(interaction)
+    
+    # Test if the output is a list
+    assert isinstance(waterid_barcode, list)
+    
+    # Test the expected waterid barcode for the sample dataframe and interaction
+    expected_waterid_barcode = [0, 104, 105]
+    assert waterid_barcode == expected_waterid_barcode