def vis_data(tmpdir_factory):
tmpdir_l = tmpdir_factory.mktemp("freqsplit_lower")
tmpdir_h = tmpdir_factory.mktemp("freqsplit_higher")
fakevis_buffer = runner.FakeVisBuffer(
num_frames=params['total_frames'],
mode=params['mode'],
freq_ids=params['freq_ids'],
wait=False
)
dump_buffer_l = runner.DumpVisBuffer(
str(tmpdir_l))
dump_buffer_h = runner.DumpVisBuffer(
str(tmpdir_h))
test = runner.KotekanStageTester(
'freqSplit', {},
fakevis_buffer,
[dump_buffer_l, dump_buffer_h],
params
)
test.run()
yield [dump_buffer_l.load(), dump_buffer_h.load()]
def run_flagging(tmpdir_factory, cmds):
""" Receive Flags """
tmpdir = tmpdir_factory.mktemp("receiveflags")
fakevis_buffer = runner.FakeVisBuffer(
num_frames=global_params["total_frames"],
mode=global_params["fakevis_mode"],
cadence=global_params["cadence"],
wait=global_params["wait"],
)
out_dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester(
"ReceiveFlags",
params,
buffers_in=fakevis_buffer,
buffers_out=out_dump_buffer,
global_config=global_params,
rest_commands=cmds,
expect_failure=True,
)
test.run()
return out_dump_buffer.load()
def written_data_base(outdir, stage_extra=None, root_extra=None):
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=writer_params["freq"],
num_frames=writer_params["total_frames"],
cadence=writer_params["cadence"],
)
root_params = writer_params.copy()
root_params["root_path"] = outdir
if root_extra is not None:
root_params.update(root_extra)
stage_params = {"node_mode": False}
if stage_extra is not None:
stage_params.update(stage_extra)
test = runner.KotekanStageTester(
"VisWriter", stage_params, fakevis_buffer, None, root_params
)
test.run()
import glob
files = sorted(glob.glob(outdir + "/20??????T??????Z_*_corr/*.h5"))
return [h5py.File(fname, "r") for fname in files]
def write_data(tmpdir_factory):
tmpdir_l = tmpdir_factory.mktemp("freqsplit_write_lower")
tmpdir_h = tmpdir_factory.mktemp("freqsplit_write_higher")
fakevis_buffer = runner.FakeVisBuffer(
num_frames=params['total_frames'],
mode=params['mode'],
freq_ids=params['freq_ids'],
wait=False
)
write_buffer_l = runner.VisWriterBuffer(
str(tmpdir_l), 'raw')
write_buffer_h = runner.VisWriterBuffer(
str(tmpdir_h), 'raw')
test = runner.KotekanStageTester(
'freqSplit', {},
fakevis_buffer,
[write_buffer_l, write_buffer_h],
params
)
test.run()
return [write_buffer_l.load(), write_buffer_h.load()]
def written_data(tmpdir_factory):
tmpdir = str(tmpdir_factory.mktemp("writer"))
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=writer_params['freq'],
num_frames=writer_params['total_frames'],
cadence=writer_params['cadence']
)
params = writer_params.copy()
params['root_path'] = tmpdir
test = runner.KotekanStageTester(
'visWriter',
{'node_mode': False, 'file_type': 'raw'},
fakevis_buffer,
None,
params
)
test.run()
import glob
files = sorted(glob.glob(tmpdir + '/20??????T??????Z_*_corr/*.meta'))
yield [visbuffer.VisRaw(fname) for fname in files]
def written_data_base(outdir, stage_extra=None, root_extra=None):
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=writer_params['freq'],
num_frames=writer_params['total_frames'],
cadence=writer_params['cadence'])
root_params = writer_params.copy()
root_params['root_path'] = outdir
if root_extra is not None:
root_params.update(root_extra)
stage_params = {
'node_mode': False,
}
if stage_extra is not None:
stage_params.update(stage_extra)
test = runner.KotekanStageTester('visWriter', stage_params, fakevis_buffer,
None, root_params)
test.run()
import glob
files = sorted(glob.glob(outdir + '/20??????T??????Z_*_corr/*.h5'))
return [h5py.File(fname, 'r') for fname in files]
def apply_data(request, tmp_path_factory, gain_path, old_gains, new_gains,
cal_broker):
output_dir = str(tmp_path_factory.mktemp("output"))
global_params["gains_dir"] = str(gain_path)
global_params[
"read_from_file"] = True if request.param == "file" else False
# REST commands
cmds = [[
"post",
"gains",
{
"update_id": new_update_id,
"start_time": new_timestamp,
"transition_interval": transition_interval,
"new_state": new_state,
},
]]
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=global_params["freq_ids"],
num_frames=global_params["total_frames"],
wait=global_params["wait"],
)
out_dump_buffer = runner.DumpVisBuffer(output_dir)
# Configuration for the direct dump of FakeVis's output
fakevis_dump_conf = {
"file_name": "fakevis_dump",
"file_ext": "dump",
"base_dir": output_dir,
}
host, port = cal_broker.server_address
global_params.update({"broker_host": host, "broker_port": port})
test = runner.KotekanStageTester(
"applyGains",
global_params,
buffers_in=fakevis_buffer,
buffers_out=out_dump_buffer,
global_config=global_params,
rest_commands=cmds,
parallel_stage_type="rawFileWrite",
parallel_stage_config=fakevis_dump_conf,
)
test.run()
in_dump = visbuffer.VisBuffer.load_files(
f"{output_dir}/*fakevis_dump*.dump")
return in_dump, out_dump_buffer.load()
def test_send_receive(tmpdir_factory):
# Run kotekan bufferRecv
tmpdir = tmpdir_factory.mktemp("writer")
write_buffer = runner.DumpVisBuffer(str(tmpdir))
# the plan is: wait 5s and then kill it
rest_commands = [("wait", 5, None), ("get", "kill", None)]
receiver = runner.KotekanStageTester(
"bufferRecv",
{},
None,
write_buffer,
params_kotekan,
rest_commands=rest_commands,
)
# TODO: network buffer processes should use in_buf and out_buf to please the test framework
receiver._stages["bufferRecv_test"]["buf"] = receiver._stages[
"bufferRecv_test"]["out_buf"]
# Run kotekan bufferRecv in another thread
with concurrent.futures.ThreadPoolExecutor() as executor:
future_receiver = executor.submit(receiver.run)
# Wait for it to start so the sender doesn't drop frames
time.sleep(1)
fakevis_buffer = runner.FakeVisBuffer(
num_frames=params_kotekan["total_frames"],
mode=params_kotekan["mode"],
freq_ids=params_kotekan["freq_ids"],
sleep_before=2,
wait=False,
)
sender = runner.KotekanStageTester("bufferSend", {}, fakevis_buffer,
None, params_kotekan)
# TODO: network buffer processes should use in_buf and out_buf to please the test framework
sender._stages["bufferSend_test"]["buf"] = sender._stages[
"bufferSend_test"]["in_buf"]
# run kotekan bufferSend
sender.run()
# wait for kotekan bufferRecv to finish
future_receiver.result(timeout=7)
assert sender.return_code == 0
assert receiver.return_code == 0
vis_data = write_buffer.load()
assert len(vis_data) == params_kotekan["total_frames"]
def data(tmpdir_factory):
# keep all the data this test produces in a tmp directory
tmpdir = tmpdir_factory.mktemp("name_of_the_test_case")
# you can use FakeVisBuffer to produce fake data
fakevis_buffer = runner.FakeVisBuffer(num_frames=params["total_frames"], mode="gaussian")
# DumpVisBuffer can be used to dump data for testing
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
# KotekanStageTester is used to run kotekan with your config
test = runner.KotekanStageTester("stageUnderTest", {}, fakevis_buffer, dump_buffer, params)
def mergewait_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("mergewait")
fakevis_fast = runner.FakeVisBuffer(freq_ids=[0],
cadence=0.3,
wait=True,
num_frames=5)
fakevis_slow = runner.FakeVisBuffer(freq_ids=[1],
cadence=5.0,
wait=True,
num_frames=10)
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester('bufferMerge', {},
[fakevis_fast, fakevis_slow], dump_buffer,
merge_params)
test.run()
yield dump_buffer.load()
def remove_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("remove")
fakevis_buffer = runner.FakeVisBuffer(
freq=remove_params['freq'], num_frames=remove_params['total_frames'])
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester('removeEv', {}, fakevis_buffer,
dump_buffer, remove_params)
test.run()
yield dump_buffer.load()
def replace_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("replace")
fakevis_buffer = runner.FakeVisBuffer(
freq=replace_params["freq"], num_frames=replace_params["total_frames"])
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester("ReplaceVis", {}, fakevis_buffer,
dump_buffer, replace_params)
test.run()
yield dump_buffer.load()
def subset_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("subset")
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=subset_params["freq_ids"],
num_frames=subset_params["total_frames"])
write_buffer = runner.VisWriterBuffer(str(tmpdir), "raw")
test = runner.KotekanStageTester("prodSubset", vis_params, fakevis_buffer,
write_buffer, subset_params)
test.run()
return write_buffer.load()
def write_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("freqsub_write")
fakevis_buffer = runner.FakeVisBuffer(num_frames=params['total_frames'],
mode=params['mode'],
freq_ids=params['freq_ids'],
wait=False)
write_buffer = runner.VisWriterBuffer(str(tmpdir), 'raw')
test = runner.KotekanStageTester('freqSubset', {}, fakevis_buffer,
write_buffer, params)
test.run()
return write_buffer.load()
def vis_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("freqsub")
fakevis_buffer = runner.FakeVisBuffer(num_frames=params['total_frames'],
mode=params['mode'],
freq_ids=params['freq_ids'],
wait=False)
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester('freqSubset', {}, fakevis_buffer,
dump_buffer, params)
test.run()
yield dump_buffer.load()
def merge_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("merge")
fakevis_buffers = [
runner.FakeVisBuffer(freq_ids=[f],
num_frames=merge_params['total_frames'])
for f in merge_params['freq']
]
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester('bufferMerge', {}, fakevis_buffers,
dump_buffer, merge_params)
test.run()
yield dump_buffer.load()
def run_eigenvis(tdir_factory, params=None):
if not params:
params = default_params
tmpdir = tdir_factory.mktemp("eigenvis")
fakevis_buffer = runner.FakeVisBuffer(freq_ids=params["freq"],
num_frames=params["total_frames"],
mode=params["mode"])
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester("eigenVis", {}, fakevis_buffer,
dump_buffer, params)
test.run()
return dump_buffer.load()
def vis_data(tmpdir_factory, request):
global num_frames
# keeping all the data this test produced here (probably do not need it)
# using FakeVisBuffer to produce fake data
fakevis_buffer = runner.FakeVisBuffer(**request.param)
num_frames = request.param["num_frames"]
# KotekanStageTester is used to run kotekan with my config
test = runner.KotekanStageTester(
stage_type="VisSharedMemWriter",
stage_config=params_writer_stage,
buffers_in=fakevis_buffer,
buffers_out=None,
global_config=global_params,
)
test.run()
def vis_data(tmpdir_factory, comet_broker):
# keeping all the data this test produced here (probably do not need it)
# using FakeVisBuffer to produce fake data
fakevis_buffer = runner.FakeVisBuffer(**params_fakevis)
# pass comet port to kotekan
params["dataset_manager"]["ds_broker_port"] = comet_broker
# KotekanStageTester is used to run kotekan with my config
test = runner.KotekanStageTester(
stage_type="VisSharedMemWriter",
stage_config=params_writer_stage,
buffers_in=fakevis_buffer,
buffers_out=None,
global_config=params,
)
yield test
def write_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("freqsub_write")
fakevis_buffer = runner.FakeVisBuffer(
num_frames=params["total_frames"],
mode=params["mode"],
freq_ids=params["freq_ids"],
wait=False,
)
write_buffer = runner.VisWriterBuffer(str(tmpdir), "raw")
test = runner.KotekanStageTester("VisFreqSubset", {}, fakevis_buffer,
write_buffer, params)
test.run()
return write_buffer.load()
def vis_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("vis_data")
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester(
'timeDownsample', downsamp_params,
runner.FakeVisBuffer(
num_frames=downsamp_params['total_frames'],
mode=downsamp_params['fakevis_mode'],
cadence=downsamp_params['cadence']
),
dump_buffer,
downsamp_params
)
test.run()
yield dump_buffer.load()
def complete_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("no_drop")
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=global_params['freq_ids'],
num_frames=global_params['total_frames'])
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
global_params.update(dump_buffer.buffer_block)
global_params.update(dump_buffer.stage_block)
global_params['compress']['out_buf'] = dump_buffer.name
valve_params = {'out_buf': 'compress_buffer'}
test = runner.KotekanStageTester('Valve', valve_params, fakevis_buffer,
None, global_params)
test.run()
yield dump_buffer.load()
def apply_data(cmds, tmpdir_factory):
apply_dir = tmpdir_factory.mktemp("apply")
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=global_params['freq_ids'],
num_frames=global_params['total_frames'],
wait=global_params['wait'])
out_dump_buffer = runner.DumpVisBuffer(str(apply_dir))
test = runner.KotekanStageTester('applyGains',
global_params,
buffers_in=fakevis_buffer,
buffers_out=out_dump_buffer,
global_config=global_params,
rest_commands=cmds)
test.run()
return out_dump_buffer.load()
def subset_data(request, tmpdir_factory):
inputs = request.param
tmpdir = tmpdir_factory.mktemp("subset")
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=subset_params["freq_ids"],
num_frames=subset_params["total_frames"])
write_buffer = runner.VisWriterBuffer(str(tmpdir), "raw")
task_params = {"inputs": inputs}
task_params.update(subset_params)
test = runner.KotekanStageTester("InputSubset", vis_params, fakevis_buffer,
write_buffer, task_params)
test.run()
return inputs, write_buffer.load()
def complete_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("no_drop")
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=global_params["freq_ids"],
num_frames=global_params["total_frames"])
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
global_params.update(dump_buffer.buffer_block)
global_params.update(dump_buffer.stage_block)
global_params["compress"]["out_buf"] = dump_buffer.name
valve_params = {"out_buf": "compress_buffer"}
test = runner.KotekanStageTester("Valve", valve_params, fakevis_buffer,
None, global_params)
test.run()
yield dump_buffer.load()
def run_flagging(tmpdir_factory, cmds):
""" Receive Flags """
tmpdir = tmpdir_factory.mktemp("receiveflags")
fakevis_buffer = runner.FakeVisBuffer(num_frames=params['total_frames'],
mode=params['fakevis_mode'],
cadence=params['cadence'],
wait=params['wait'])
out_dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester('receiveFlags',
params,
buffers_in=fakevis_buffer,
buffers_out=out_dump_buffer,
global_config=params,
rest_commands=cmds)
test.run()
return out_dump_buffer.load()
def diagonal_data(tmpdir_factory):
tmpdir = tmpdir_factory.mktemp("diagonal")
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=diag_global_params["freq_ids"],
num_frames=diag_global_params["total_frames"],
)
dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester(
"baselineCompression",
diag_stage_params,
fakevis_buffer,
dump_buffer,
diag_global_params,
)
test.run()
yield dump_buffer.load()
def vis_data_zero_weights(tmpdir_factory):
""" Truncated visibilities """
tmpdir = tmpdir_factory.mktemp("vis_data_t")
fakevis_buffer = runner.FakeVisBuffer(
num_frames=trunc_params["total_frames"],
mode=trunc_params["fakevis_mode"],
cadence=trunc_params["cadence"],
zero_weight=True,
)
in_dump_config = trunc_params.copy()
in_dump_config["base_dir"] = str(tmpdir)
in_dump_config["file_name"] = "fakevis"
in_dump_config["file_ext"] = "dump"
out_dump_buffer = runner.DumpVisBuffer(str(tmpdir))
test = runner.KotekanStageTester(
"VisTruncate",
trunc_params,
buffers_in=fakevis_buffer,
buffers_out=out_dump_buffer,
global_config=trunc_params,
parallel_stage_type="rawFileWrite",
parallel_stage_config=in_dump_config,
noise="random",
)
test.run()
yield (
out_dump_buffer.load(),
visbuffer.VisBuffer.load_files("%s/*fakevis*.dump" % str(dir)),
)
def transpose(tmpdir_factory, cal_broker):
writer_params["file_length"] = writer_params["total_frames"]
# Write fake data in raw format
tmpdir = str(tmpdir_factory.mktemp("writer"))
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=writer_params["freq"],
num_frames=writer_params["total_frames"],
cadence=writer_params["cadence"],
mode=writer_params["mode"],
test_pattern_value=writer_params["test_pattern_value"],
wait=True, # make sure cal_broker has time to update
)
# Add an applyGain stage to alter the dataset ID
gain_buf_name = "gains_applied"
fakevis_buffer.buffer_block.update({
gain_buf_name: {
"kotekan_buffer": "vis",
"metadata_pool": "vis_pool",
"num_frames": "buffer_depth",
},
})
host, port = cal_broker.server_address
fakevis_buffer.stage_block.update({
"apply_gains": {
"kotekan_stage": "applyGains",
"in_buf": fakevis_buffer.name,
"out_buf": gain_buf_name,
"log_level": "debug",
"broker_host": host,
"broker_port": port,
},
})
# Remove samples from two frequency to test handling of empty frames
fsel_frac_name = "frac_freqsel"
fsel_empty_name = "empty_freqsel"
fsel_buf = {
fsel_frac_name: {
"kotekan_buffer": "vis",
"metadata_pool": "vis_pool",
"num_frames": "buffer_depth",
},
fsel_empty_name: {
"kotekan_buffer": "vis",
"metadata_pool": "vis_pool",
"num_frames": "buffer_depth",
},
}
fakevis_buffer.buffer_block.update(fsel_buf)
fakevis_buffer.stage_block.update({
"frac_fsel": {
"kotekan_stage": "visDrop",
"in_buf": gain_buf_name,
"out_buf": fsel_frac_name,
"freq": [writer_params["freq"][frac_freq]],
"frac_lost": frac_lost,
"frac_rfi": frac_rfi,
"log_level": "debug",
},
"empty_fsel": {
"kotekan_stage": "visDrop",
"in_buf": fsel_frac_name,
"out_buf": fsel_empty_name,
"freq": [writer_params["freq"][empty_freq]],
"log_level": "debug",
},
})
# update the output buffer name
fakevis_buffer.name = fsel_empty_name
# REST commands for gains update
cmds = [
["wait", 2.0, {}],
[
"post",
"gains",
{
"update_id": new_update_id,
"start_time": new_timestamp,
"transition_interval": transition_interval,
"new_state": new_state,
},
],
]
# Write out the test data in raw and HDF5 simultaneously
tmpdir_h5 = str(tmpdir_factory.mktemp("dump_h5"))
dumph5_conf = writer_params.copy()
dumph5_conf["root_path"] = str(tmpdir_h5)
dumph5_conf["file_name"] = "dumph5"
dumph5_conf["node_mode"] = False
params = writer_params.copy()
params["root_path"] = tmpdir
writer = runner.KotekanStageTester(
"VisWriter",
{
"node_mode": False,
"write_ev": True,
"file_type": "raw"
},
fakevis_buffer,
None,
params,
parallel_stage_type="VisWriter",
parallel_stage_config=dumph5_conf,
noise="random",
rest_commands=cmds,
)
writer.run()
# get raw infile
files = sorted(glob.glob(tmpdir + "/20??????T??????Z_*_corr/*.meta"))
assert len(files) == 1
infile = os.path.splitext(files[0])[0]
# get hdf5 infile
files = sorted(glob.glob(tmpdir_h5 + "/20??????T??????Z_*_corr/*.h5"))
assert len(files) == 1
infile_h5 = os.path.splitext(files[0])[0]
# Tranpose and write data
raw_buf = runner.ReadRawBuffer(infile, writer_params["chunk_size"])
outfile = tmpdir + "/transposed"
transposer = runner.KotekanStageTester(
"VisTranspose",
{
"outfile": outfile,
"chunk_size": writer_params["chunk_size"],
"comet_timeout": 120.0,
},
raw_buf,
None,
params,
)
transposer.run()
fh = h5py.File(infile_h5 + ".h5", "r")
fh_t = h5py.File(outfile + ".h5", "r")
yield (fh_t, fh)
fh.close()
fh_t.close()
def transpose_stack(tmpdir_factory):
# Write fake stacked data in raw format
tmpdir = str(tmpdir_factory.mktemp("writer"))
fakevis_buffer = runner.FakeVisBuffer(
freq_ids=stack_params["freq"],
num_frames=stack_params["file_length"],
cadence=stack_params["cadence"],
mode="chime",
)
# Add stacking stage
stack_buf_name = "fake_stacked"
stack_buf = {
stack_buf_name: {
"kotekan_buffer": "vis",
"metadata_pool": "vis_pool",
"num_frames": "buffer_depth",
}
}
fakevis_buffer.buffer_block.update(stack_buf)
fakevis_buffer.stage_block.update({
"fakevis_stack": {
"kotekan_stage": "baselineCompression",
"in_buf": fakevis_buffer.name,
"out_buf": stack_buf_name,
"stack_type": "chime_in_cyl",
}
})
fakevis_buffer.name = stack_buf_name
params = stack_params.copy()
params["root_path"] = tmpdir
writer = runner.KotekanStageTester(
"VisWriter",
{
"node_mode": False,
"write_ev": True,
"file_type": "raw"
},
fakevis_buffer,
None,
params,
)
writer.run()
# get raw infile
files = sorted(glob.glob(tmpdir + "/20??????T??????Z_*_corr/*.meta"))
assert len(files) == 1
infile = os.path.splitext(files[0])[0]
# Tranpose and write data
raw_buf = runner.ReadRawBuffer(infile, stack_params["chunk_size"])
outfile = tmpdir + "/transposed"
transposer = runner.KotekanStageTester(
"VisTranspose",
{
"outfile": outfile,
"infile": infile,
"chunk_size": writer_params["chunk_size"],
"comet_timeout": 120.0,
},
raw_buf,
None,
params,
)
transposer.run()
fh = h5py.File(outfile + ".h5", "r")
yield (infile, fh)
fh.close()
