From 8ed9b299366a307b996b5eff4124ee47998f3690 Mon Sep 17 00:00:00 2001
From: Zain Sohail <zainsohail04@gmail.com>
Date: Mon, 6 Nov 2023 14:20:40 +0100
Subject: [PATCH] allow to test all formats of channels and if group_name is
 not in h5

---
 tests/data/loader/flash/config.yaml     |  4 ++
 tests/loader/flash/test_flash_loader.py | 62 ++++++++++++++++++++++++-
 2 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/tests/data/loader/flash/config.yaml b/tests/data/loader/flash/config.yaml
index 6989f72d..d04bd56f 100644
--- a/tests/data/loader/flash/config.yaml
+++ b/tests/data/loader/flash/config.yaml
@@ -118,6 +118,10 @@ dataframe:
       format: per_train
       group_name: "/uncategorised/FLASH.DIAG/TIMINGINFO/TIME1.BUNCH_FIRST_INDEX.1/"
 
+    delayStage:
+      format: per_pulse
+      group_name: "/zraw/FLASH.SYNC/LASER.LOCK.EXP/F1.PG.OSC/FMC0.MD22.1.ENCODER_POSITION.RD/dGroup"
+
   # The prefixes of the stream names for different DAQ systems for parsing filenames
   # (Not to be changed by user)
   stream_name_prefixes:
diff --git a/tests/loader/flash/test_flash_loader.py b/tests/loader/flash/test_flash_loader.py
index 84119aa0..c2f044a4 100644
--- a/tests/loader/flash/test_flash_loader.py
+++ b/tests/loader/flash/test_flash_loader.py
@@ -1,6 +1,7 @@
 import os
 from importlib.util import find_spec
 from pathlib import Path
+from typing import Literal
 
 import pytest
 
@@ -16,7 +17,7 @@ def config_file():
     return parse_config(config_path)
 
 
-def test_get_channels_by_format(config_file):
+def test_get_channels_by_format(config_file: dict):
 
     fl = FlashLoader(config_file)
     electron_channels = ["dldPosX", "dldPosY", "dldTimeSteps"]
@@ -28,15 +29,19 @@ def test_get_channels_by_format(config_file):
         "cryoTemperature",
         "sampleTemperature",
         "dldTimeBinSize",
+        "delayStage",
     ]
+    train_channels = ["timeStamp"]
 
     # Call get_channels_by_format method
     format_electron = fl.get_channels_by_format(["per_electron"])
     format_pulse = fl.get_channels_by_format(["per_pulse"])
+    format_train = fl.get_channels_by_format(["per_train"])
     format_both = fl.get_channels_by_format(["per_pulse", "per_electron"])
 
     assert set(electron_channels) == set(format_electron)
     assert set(pulse_channels) == set(format_pulse)
+    assert set(train_channels) == set(format_train)
     assert set(electron_channels + pulse_channels) == set(format_both)
 
 
@@ -44,7 +49,11 @@ def test_get_channels_by_format(config_file):
     "sub_dir",
     ["online-0/fl1user3/", "express-0/fl1user3/", "FL1USER3/"],
 )
-def test_initialize_paths(config_file, fs, sub_dir):
+def test_initialize_paths(
+    config_file: dict,
+    fs,
+    sub_dir: Literal["online-0/fl1user3/", "express-0/fl1user3/", "FL1USER3/"],
+):
     config = config_file
     del config["core"]["paths"]
     config["core"]["beamtime_id"] = "12345678"
@@ -69,3 +78,52 @@ def test_initialize_paths(config_file, fs, sub_dir):
     print(data_raw_dir)
     assert expected_raw_path == data_raw_dir[0]
     assert expected_processed_path == data_parquet_dir
+
+
+def test_initialize_paths_filenotfound(config_file: dict):
+    # test the FileNotFoundError
+    config = config_file
+    del config["core"]["paths"]
+    config["core"]["beamtime_id"] = "11111111"
+    config["core"]["year"] = "2000"
+
+    # instance of class with correct config and call initialize_paths
+    fl = FlashLoader(config=config)
+    with pytest.raises(FileNotFoundError):
+        data_raw_dir, data_parquet_dir = fl.initialize_paths()
+
+
+def test_invalid_channel_format(config_file: dict):
+    config = config_file
+    config["dataframe"]["channels"]["dldPosX"]["format"] = "foo"
+
+    fl = FlashLoader(config=config)
+
+    with pytest.raises(ValueError):
+        fl.read_dataframe()
+
+
+def test_group_name_not_in_h5(config_file: dict):
+    config = config_file
+    config["dataframe"]["channels"]["dldPosX"]["group_name"] = "foo"
+    h5_path = "FLASH1_USER3_stream_2_run43878_file1_20230130T153807.1.h5"
+    fl = FlashLoader(config=config)
+
+    with pytest.raises(ValueError) as e:
+        fl.create_dataframe_per_file(config["core"]["paths"]["data_raw_dir"] + h5_path)
+        print(e)
+    assert str(e.value.args[0]) == "The group_name for channel dldPosX does not exist."
+
+
+# def test_buffer_schema_mismatch(config_file: dict):
+#     fl = FlashLoader(config=config_file)
+
+#     fl.read_dataframe(runs=[43878])
+
+#     config = config_file
+#     config["dataframe"]["channels"]["dldPosX2"] = {
+#         "group_name": "/uncategorised/FLASH.EXP/HEXTOF.DAQ/DLD1/"}
+
+#     fl = FlashLoader(config=config)
+#     with pytest.raises(ValueError):
+#         fl.read_dataframe(runs=[43878])