text: import os
  import shutil
  import unittest
  
  import gwpy
  import mock
  
  import bilby
  from bilby_pipe.data_generation import DataGenerationInput, create_generation_parser
  from bilby_pipe.main import parse_args
  from bilby_pipe.utils import BilbyPipeError, DataDump
  
  
  class TestDataGenerationInput(unittest.TestCase):
      def setUp(self):
          self.outdir = "test_outdir"
          self.default_args_list = [
              "--ini",
              "tests/test_data_generation.ini",
              "--outdir",
              self.outdir,
              "--data-label",
              "TEST",
          ]
          self.parser = create_generation_parser()
          self.inputs = DataGenerationInput(
              *parse_args(self.default_args_list, self.parser), create_data=False
          )
          self.gps_file = "tests/gps_file.txt"
  
      def tearDown(self):
          del self.default_args_list
          del self.parser
          del self.inputs
          if os.path.isdir(self.outdir):
              shutil.rmtree(self.outdir)
  
      def test_cluster_set(self):
          self.inputs.cluster = 123
          self.assertEqual(123, self.inputs.cluster)
  
      def test_process_set(self):
          self.inputs.process = 321
          self.assertEqual(321, self.inputs.process)
  
      def test_parameter_conversion(self):
          self.inputs.frequency_domain_source_model = "binary_neutron_star"
          self.assertEqual(
              self.inputs.parameter_conversion,
              bilby.gw.conversion.convert_to_lal_binary_neutron_star_parameters,
          )
          self.inputs.frequency_domain_source_model = "binary_black_hole"
          self.assertEqual(
              self.inputs.parameter_conversion,
              bilby.gw.conversion.convert_to_lal_binary_black_hole_parameters,
          )
  
      def test_psd_length_set(self):
          self.inputs.psd_length = 10
          self.assertEqual(10, self.inputs.psd_length)
          self.assertEqual(
              10 * self.inputs.duration, self.inputs.psd_length * self.inputs.duration
          )
  
      def test_psd_length(self):
          self.inputs.duration = 4
          self.inputs.psd_length = 32
          self.assertEqual(32, self.inputs.psd_length)
          self.assertEqual(4, self.inputs.duration)
          self.assertEqual(
              self.inputs.psd_length * self.inputs.duration, self.inputs.psd_duration
          )
  
      def test_psd_length_over_maximum(self):
          self.inputs.duration = 128
          self.inputs.psd_length = 32
          self.assertEqual(32, self.inputs.psd_length)
          self.assertEqual(128, self.inputs.duration)
          self.assertEqual(1024, self.inputs.psd_duration)
  
      def test_psd_length_over_custom_maximum(self):
          self.inputs.psd_maximum_duration = 2048
          self.inputs.duration = 128
          self.inputs.psd_length = 32
          self.assertEqual(32, self.inputs.psd_length)
          self.assertEqual(128, self.inputs.duration)
          self.assertEqual(2048, self.inputs.psd_duration)
  
      def test_set_reference_frequency(self):
          args_list = self.default_args_list + ["--reference-frequency", "10"]
          inputs = DataGenerationInput(
              *parse_args(args_list, self.parser), create_data=False
          )
          self.assertEqual(inputs.reference_frequency, 10)
  
      def test_psd_length_default(self):
          self.assertEqual(32 * self.inputs.duration, self.inputs.psd_duration)
  
      def test_psd_start_time_set(self):
          self.inputs.psd_start_time = 10
          self.assertEqual(10, self.inputs.psd_start_time)
  
      def test_psd_start_time_default(self):
          self.inputs.psd_duration = 4
          self.inputs.trigger_time = 12
          self.assertEqual(-4, self.inputs.psd_start_time)
  
      # def test_psd_start_time_fail(self):
      #    self.inputs.psd_duration = 4
      #    self.inputs.duration = 4
      #    self.inputs.trigger_time = 2
      #    self.inputs.start_time = 0
      #    self.inputs.psd_start_time = None
      #    with self.assertRaises(BilbyPipeError):
      #        self.assertEqual(10 - 4, self.inputs.psd_start_time)
  
      def test_script_inputs_from_test_ini(self):
          self.assertEqual(
              self.inputs.channel_dict, dict(H1="GDS-CALIB_STRAIN", L1="GDS-CALIB_STRAIN")
          )
          self.assertEqual(self.inputs.label, "label")
  
      def test_interferometer_unset(self):
          with self.assertRaises(ValueError):
              self.inputs.interferometers
  
      def test_interferometer_set(self):
          ifos = bilby.gw.detector.InterferometerList(["H1", "L1"])
          self.inputs.interferometers = ifos
          self.assertEqual(ifos, self.inputs.interferometers)
  
      def test_script_inputs_detectors_from_ini(self):
          self.assertEqual(self.inputs.detectors, ["H1", "L1"])
  
      def test_script_inputs_detectors_from_command_line(self):
          args_list = self.default_args_list + ["--detectors", "H1", "--detectors", "L1"]
          inputs = DataGenerationInput(
              *parse_args(args_list, self.parser), create_data=False
          )
          self.assertEqual(inputs.detectors, ["H1", "L1"])
  
          args_list = self.default_args_list + ["--detectors", "H1 L1"]
          inputs = DataGenerationInput(
              *parse_args(args_list, self.parser), create_data=False
          )
          self.assertEqual(inputs.detectors, ["H1", "L1"])
  
          args_list = self.default_args_list + ["--detectors", "L1 H1"]
          inputs = DataGenerationInput(
              *parse_args(args_list, self.parser), create_data=False
          )
          self.assertEqual(inputs.detectors, ["H1", "L1"])
  
          args_list = self.default_args_list + ["--detectors", "[L1, H1]"]
          inputs = DataGenerationInput(
              *parse_args(args_list, self.parser), create_data=False
          )
  
          args_list = self.default_args_list + ["--detectors", "[L1 H1]"]
          inputs = DataGenerationInput(
              *parse_args(args_list, self.parser), create_data=False
          )
          self.assertEqual(inputs.detectors, ["H1", "L1"])
  
          args_list = self.default_args_list + ["--detectors", '["L1", "H1"]']
          inputs = DataGenerationInput(
              *parse_args(args_list, self.parser), create_data=False
          )
          self.assertEqual(inputs.detectors, ["H1", "L1"])
  
          args_list = self.default_args_list + ["--detectors", "['L1', 'H1']"]
          inputs = DataGenerationInput(
              *parse_args(args_list, self.parser), create_data=False
          )
          self.assertEqual(inputs.detectors, ["H1", "L1"])
  
      def test_detectors_not_understood(self):
          with self.assertRaises(BilbyPipeError):
              self.inputs.detectors = 10
  
      # def test_trigger_time(self):
      #    args_list = [
      #        "--ini",
      #        "tests/test_data_generation.ini",
      #        "--outdir",
      #        self.outdir,
      #        "--trigger-time",
      #        "1126259462",
      #        "--data-label",
      #        "TEST",
      #    ]
      #    self.inputs = DataGenerationInput(*parse_args(args_list, self.parser))
  
      def test_injections_no_file(self):
          args_list = [
              "--ini",
              "tests/test_data_generation.ini",
              "--outdir",
              self.outdir,
              "--injection-file",
              "not_a_file",
              "--data-label",
              "TEST",
          ]
          with self.assertRaises(FileNotFoundError):
              self.inputs = DataGenerationInput(*parse_args(args_list, self.parser))
  
      def test_generation_seed_increases_with_injection_index(self):
          """Assert that generation seed increments for each job.
  
          ie
          JOB 0 -- seed X
          JOB 1 -- seed X + 1
          ...
          JOB N -- seed X + N
  
          This is so that the gaussian data for each job will be different.
          """
          idx = 0
          generation_seed = 0
          args_list = [
              "--ini",
              "tests/test_data_generation.ini",
              f"--generation-seed={generation_seed}",
              f"--idx={idx}",
              "--gaussian-noise",
              "--trigger-time",
              "2",
              "--outdir",
              self.outdir,
              "--label",
              "TEST",
          ]
          self.inputs = DataGenerationInput(*parse_args(args_list, self.parser))
          self.assertEqual(self.inputs.generation_seed, idx + generation_seed)
          idx = 2
          generation_seed = 0
          args_list = [
              "--ini",
              "tests/test_data_generation.ini",
              f"--generation-seed={generation_seed}",
              f"--idx={idx}",
              "--gaussian-noise",
              "--trigger-time",
              "1126259462",
              "--outdir",
              self.outdir,
              "--label",
              "TEST",
          ]
          self.inputs = DataGenerationInput(*parse_args(args_list, self.parser))
          self.assertEqual(self.inputs.generation_seed, idx + generation_seed)
  
      def test_generation_seed_is_random_if_none_provided(self):
          """Assert that the generation seed is some random value if not provided."""
          idx = 0
          generation_seed = None
          args_list = [
              "--ini",
              "tests/test_data_generation.ini",
              f"--generation-seed={generation_seed}",
              f"--idx={idx}",
              "--gaussian-noise",
              "--trigger-time",
              "1126259462",
              "--outdir",
              self.outdir,
              "--label",
              "TEST",
          ]
          self.inputs = DataGenerationInput(*parse_args(args_list, self.parser))
          self.assertTrue(1 <= self.inputs.generation_seed <= 1e6)
  
      @mock.patch("bilby_pipe.data_generation.DataGenerationInput._get_data")
      @mock.patch("bilby.gw.detector.inject_signal_into_gwpy_timeseries")
      def test_inject_signal_into_time_domain_data(
          self, inject_signal_into_timeseries_method, get_data_method
      ):
          timeseries, metadata = load_test_strain_data()
  
          get_data_method.return_value = timeseries
          inject_signal_into_timeseries_method.return_value = (timeseries, metadata)
          args_list = ["tests/test_injection.ini", "--outdir", self.outdir]
  
          inputs = DataGenerationInput(*parse_args(args_list, self.parser))
          self.assertTrue(inputs.injection_parameters["geocent_time"] == 0)
          self.assertEqual(inject_signal_into_timeseries_method.call_count, 2)
          self.assertTrue(get_data_method.called)
  
          t0 = 1126259463.4
          t1 = t0 + 1
  
          t0_psd = t0 - 32
          t1_psd = t0
  
          get_data_method.assert_any_call("H1", "GWOSC", t0, t1)  # SIGNAL
          get_data_method.assert_any_call("H1", "GWOSC", t0_psd, t1_psd)  # PSD
          get_data_method.assert_any_call("L1", "GWOSC", t0, t1)  # SIGNAL
          get_data_method.assert_any_call("L1", "GWOSC", t0_psd, t1_psd)  # PSD
  
      def test_inject_signal_into_gaussian_noise(self):
          args_list = [
              "tests/test_injection_in_gaussian_noise.ini",
              "--outdir",
              self.outdir,
          ]
          data_input = DataGenerationInput(*parse_args(args_list, self.parser))
          injection_param = data_input.injection_parameters
          self.assertTrue(injection_param["geocent_time"] == 0)
  
      @mock.patch("bilby_pipe.data_generation.DataGenerationInput._gwpy_get")
      @mock.patch("bilby_pipe.data_generation.DataGenerationInput._is_gwpy_data_good")
      @mock.patch("bilby_pipe.data_generation.logger")
      def test_data_quality_ignore_flag(self, mock_logs, is_data_good, get_data_method):
          timeseries, _ = load_test_strain_data()
          is_data_good.return_value = False
          get_data_method.return_value = timeseries
  
          args_list = [
              "tests/test_basic_ini.ini",
              "--detectors",
              "[H1, L1]",
              "--channel-dict",
              "{'H1': 'GDS-CALIB_STRAIN', 'L1': 'GDS-CALIB_STRAIN'}",
              "--duration",
              " 1",
              "--prior_file",
              "tests/example_prior.prior",
              "--waveform-approximant",
              "IMRPhenomPv2",
              "--idx",
              "0",
              "--trigger_time",
              "1126259462.4",
              "--label",
              "QUALITY_TEST",
          ]
  
          # make sure that when the flag is present, no error
          args, unknown = parse_args(args_list, create_generation_parser())
          args.trigger_time = 1126259462.4
          input = DataGenerationInput(args, unknown)
          self.assertFalse(input._is_gwpy_data_good())
          self.assertTrue(input.ignore_gwpy_data_quality_check)
  
          # make sure that when the flag is not present, error present
          args, unknown = parse_args(args_list, create_generation_parser())
          args.trigger_time = 1126259462.4
          args.ignore_gwpy_data_quality_check = False
          with self.assertRaises(BilbyPipeError):
              DataGenerationInput(args, unknown)
              self.assertFalse(input._is_gwpy_data_good())
              self.assertFalse(input.ignore_gwpy_data_quality_check)
  
      @mock.patch("gwpy.segments.DataQualityFlag.query")
      @mock.patch("bilby_pipe.data_generation.logger")
      def test_data_quality_fail(self, mock_logs, quality_query):
          """Test the data quality check function's FAIL state.
  
          Parameters
          ----------
          mock_logs: the logging module being used inside this function
  
          """
          full_data = gwpy.segments.DataQualityFlag.read("tests/DATA/data_quality.hdf5")
          start_time_bad, end_time_bad = 1241725028.9, 1241725029.1
          quality_query.return_value = full_data
          data_is_good = DataGenerationInput._is_gwpy_data_good(
              start_time=start_time_bad, end_time=end_time_bad, det="H1"
          )
          self.assertFalse(data_is_good)
          self.assertTrue(mock_logs.warning.called)
          warning_log_str = mock_logs.warning.call_args.args[0]
          self.assertTrue("Data quality check: FAILED" in warning_log_str)
  
      @mock.patch("gwpy.segments.DataQualityFlag.query")
      @mock.patch("bilby_pipe.data_generation.logger")
      def test_data_quality_pass(self, mock_logs, quality_query):
          """Test the data quality function's PASS state.
  
          Parameters
          ----------
          mock_logs: the logging module being used inside this function
  
          """
          full_data = gwpy.segments.DataQualityFlag.read("tests/DATA/data_quality.hdf5")
          start_time_good, end_time_good = 1241725028.9, 1241725029
          quality_query.return_value = full_data
          data_is_good = DataGenerationInput._is_gwpy_data_good(
              start_time=start_time_good, end_time=end_time_good, det="H1"
          )
          self.assertTrue(data_is_good)
          self.assertFalse(mock_logs.warning.called)
  
      @mock.patch("gwpy.segments.DataQualityFlag.query")
      @mock.patch("bilby_pipe.data_generation.logger")
      def test_data_quality_exception(self, mock_logs, quality_query):
          """Test the data quality function's PASS state.
  
          Parameters
          ----------
          mock_logs: the logging module being used inside this function
  
          """
          start_time_good, end_time_good = 1241725028.9, 1241725029
          quality_query.side_effect = Exception("Some exception from GWpy")
          data_is_good = DataGenerationInput._is_gwpy_data_good(
              start_time=start_time_good, end_time=end_time_good, det="H1"
          )
          self.assertTrue(data_is_good is None)
          self.assertTrue(mock_logs.warning.called)
  
  
  class TestDataReading(unittest.TestCase):
      def setUp(self):
          self.outdir = "test_outdir"
          self.data_dir = "tests/DATA/"
          self.default_args_list = [
              "--ini",
              "tests/test_data_generation.ini",
              "--outdir",
              self.outdir,
              "--data-label",
              "TEST",
          ]
          self.parser = create_generation_parser()
          self.inputs = DataGenerationInput(
              *parse_args(self.default_args_list, self.parser), create_data=False
          )
  
          self.det = "H1"
          self.channel = "H1:DCS-CALIB_STRAIN_C02"
          self.start_time = 1126259356.0
          self.end_time = 1126259357.0
  
      def tearDown(self):
          del self.inputs
          if os.path.isdir(self.outdir):
              shutil.rmtree(self.outdir)
  
      def test_read_data_gwf(self):
          self.inputs.data_dict = {self.det: f"{self.data_dir}/test_data.gwf"}
          data = self.inputs._gwpy_read(
              self.det, self.channel, self.start_time, self.end_time
          )
          self.assertEqual(data.times[0].value, self.start_time)
          self.assertEqual(len(data), 16384)
  
      def test_read_data_txt(self):
          self.inputs.data_dict = {self.det: f"{self.data_dir}/test_data.txt"}
          data = self.inputs._gwpy_read(
              self.det, self.channel, self.start_time, self.end_time
          )
          self.assertEqual(data.times[0].value, self.start_time)
          self.assertEqual(len(data), 16384)
  
      def test_read_data_hdf5(self):
          self.inputs.data_dict = {self.det: f"{self.data_dir}/test_data.hdf5"}
          data = self.inputs._gwpy_read(
              self.det, self.channel, self.start_time, self.end_time
          )
          self.assertEqual(data.times[0].value, self.start_time)
          self.assertEqual(len(data), 16384)
  
  
  def load_test_strain_data():
      """Helper function to load data from gwpy_data.pickle"""
      ifo = DataDump.from_pickle("tests/DATA/gwpy_data.pickle").interferometers[0]
      timeseries = ifo.strain_data.to_gwpy_timeseries()
      metadata = ifo.meta_data
      return timeseries, metadata
  
  
  if __name__ == "__main__":
      unittest.main()
  