conftest.py

#
# Copyright (C) 2021 Associated Universities, Inc. Washington DC, USA.
#
# This file is part of NRAO Workspaces.
#
# Workspaces is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Workspaces is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Workspaces.  If not, see <https://www.gnu.org/licenses/>.
""" Various 'n' sundry utilities for our tests """

# pylint: disable=E0401, R1721

from os import mkdir
from pathlib import Path
from typing import List

import pytest
from ingest_envoy.manifest_components import (
    INIT_WEBLOG_FILENAME,
    TARFILE_EXT,
    WEBLOG_FILENAME,
    AncillaryProduct,
    OutputScienceProduct,
)
from ingest_envoy.utilities import AncillaryProductType

# --------------------------------
#    EVLA CAL manifest test data
# --------------------------------
EVLA_CAL_INPUT_FILENAMES = ["20A-346_2021_07_23_T13_37_08.376.tar", WEBLOG_FILENAME]
UNWANTED = ["ignore_me.fits", "just_a_lotta_nothing", "uninteresting_metadata.xml"]


@pytest.fixture(scope="function")
def ingest_path(tmp_path: Path) -> Path:
    """
    Make a directory to use as the ingestion staging dir, or curation source

    :param tmp_path: built-in pytest fixture, Pytest cleans them up periodically
    :return: Path to new directory
    """

    # cast is necessary because otherwise we get a LocalPath, which doesn't work
    fake_ingest_path = tmp_path / "ingestion"
    fake_ingest_path.mkdir()
    return fake_ingest_path


@pytest.fixture
def alternate_manifest_destination(tmp_path: Path) -> Path:
    """
    Make an alternative directory to ingest_path for tests to put their manifests in

    :param tmp_path: built-in pytest fixture, Pytest cleans them up periodically
    :return: Path to new directory
    """
    alternate_manifest_destination = tmp_path / "manifest_destination"
    alternate_manifest_destination.mkdir()
    return alternate_manifest_destination


def find_example_manifest(manifest_name: str) -> Path:
    """
    Get this example manifest for comparison with one we've generated in a test.

    :param manifest_name: unique file identifier
    :return: full path to the manifest file
    """
    filename = manifest_name + ".json"
    for file in Path.cwd().rglob(filename):
        return file

    raise FileNotFoundError(filename)


def populate_fake_evla_cal_ingest_path(staging_dir: Path) -> List[Path]:
    """
    Create a directory containing fake calibration products, plus other stuff
    that we -don't- want to ingest.

    :param staging_dir: our temporary dir
    :return:
    """

    files = []
    filenames = [filename for filename in EVLA_CAL_INPUT_FILENAMES]
    for filename in UNWANTED:
        filenames.append(filename)

    for filename in filenames:
        path = staging_dir / filename
        path.touch()
        files.append(path)

    return files


def populate_fake_final_evla_cal_ingest_path(staging_dir: Path) -> List[Path]:
    """
    Create a directory containing fake calibration products, plus other stuff
    that we -don't- want to ingest, PLUS an initial weblog

    :param staging_dir: our temporary dir
    :return:
    """
    files = populate_fake_evla_cal_ingest_path(staging_dir)
    init_weblog = staging_dir / "initial_weblog.tgz"
    init_weblog.touch()
    files.append(init_weblog)
    return files


# -----------------------------
#    Image manifest test data
# -----------------------------
IMAGE_LOCATOR = "uid://evla/calibration/ea93dae5-3495-47fa-887d-4be2852f5f14"

ADDITIONAL_METADATA_FILENAME = "aux_image_metadata.json"

PRIMARY_BEAM_ANCILLARY_K = AncillaryProduct(
    type=AncillaryProductType.PB_FITS, filename="16B-069.MJD57713.51329133102.J1522+3934_sci.K_band.cont.I.pb.tt0.fits"
)
CLEAN_MASK_ANCILLARY_K = AncillaryProduct(
    type=AncillaryProductType.FITS_MASK, filename="16B-069.MJD57713.51329133102.J1522+3934_sci.K_band.cont.I.mask.fits"
)
ALPHA_ERROR_K = AncillaryProduct(
    type=AncillaryProductType.FITS,
    filename="16B-069.MJD57713.51329133102.J1522+3934_sci.K_band.cont.I.alpha.error.fits",
)
ALPHA_K = AncillaryProduct(
    type=AncillaryProductType.FITS, filename="16B-069.MJD57713.51329133102.J1522+3934_sci.K_band.cont.I.alpha.fits"
)
TT0_K = AncillaryProduct(
    type=AncillaryProductType.FITS, filename="16B-069.MJD57713.51329133102.J1522+3934_sci.K_band.cont.I.tt0.fits"
)
TT1_K = AncillaryProduct(
    type=AncillaryProductType.FITS, filename="16B-069.MJD57713.51329133102.J1522+3934_sci.K_band.cont.I.tt1.fits"
)
OUTPUT_GROUP_SCIENCE_PRODUCT_ANCILLARIES_K = [
    PRIMARY_BEAM_ANCILLARY_K,
    CLEAN_MASK_ANCILLARY_K,
    ALPHA_ERROR_K,
    ALPHA_K,
    TT0_K,
    TT1_K,
]

OUTPUT_SCIENCE_PRODUCT_K = OutputScienceProduct(
    product_type=AncillaryProductType.FITS,
    filename="16B-069.MJD57713.51329133102.J1522+3934_sci.K_band.cont.I.pbcor.tt0.fits",
    ancillary_products=OUTPUT_GROUP_SCIENCE_PRODUCT_ANCILLARIES_K,
)

PRIMARY_BEAM_ANCILLARY_X = AncillaryProduct(
    type=AncillaryProductType.PB_FITS, filename="16B-069.MJD57713.51329133102.J1522+3934_sci.X_band.cont.I.pb.tt0.fits"
)
CLEAN_MASK_ANCILLARY_X = AncillaryProduct(
    type=AncillaryProductType.FITS_MASK, filename="16B-069.MJD57713.51329133102.J1522+3934_sci.X_band.cont.I.mask.fits"
)
ALPHA_ERROR_X = AncillaryProduct(
    type=AncillaryProductType.FITS,
    filename="16B-069.MJD57713.51329133102.J1522+3934_sci.X_band.cont.I.alpha.error.fits",
)
ALPHA_X = AncillaryProduct(
    type=AncillaryProductType.FITS, filename="16B-069.MJD57713.51329133102.J1522+3934_sci.X_band.cont.I.alpha.fits"
)
TT0_X = AncillaryProduct(
    type=AncillaryProductType.FITS, filename="16B-069.MJD57713.51329133102.J1522+3934_sci.X_band.cont.I.tt0.fits"
)
TT1_K = AncillaryProduct(
    type=AncillaryProductType.FITS, filename="16B-069.MJD57713.51329133102.J1522+3934_sci.X_band.cont.I.tt1.fits"
)
OUTPUT_GROUP_SCIENCE_PRODUCT_ANCILLARIES_X = [
    PRIMARY_BEAM_ANCILLARY_X,
    CLEAN_MASK_ANCILLARY_X,
    ALPHA_ERROR_X,
    ALPHA_X,
    TT0_X,
    TT1_K,
]

OUTPUT_SCIENCE_PRODUCT_X = OutputScienceProduct(
    product_type=AncillaryProductType.FITS,
    filename="16B-069.MJD57713.51329133102.J1522+3934_sci.X_band.cont.I.pbcor.tt0.fits",
    ancillary_products=OUTPUT_GROUP_SCIENCE_PRODUCT_ANCILLARIES_X,
)

# input files
WEBLOG_ANCILLARY = AncillaryProduct(type=AncillaryProductType.PIPELINE_WEBLOG, filename=WEBLOG_FILENAME)
INIT_WEBLOG_ANCILLARY = AncillaryProduct(type=AncillaryProductType.PIPELINE_WEBLOG, filename=INIT_WEBLOG_FILENAME)
PIPELINE_AF_ANCILLARY = AncillaryProduct(
    type=AncillaryProductType.PIPELINE_ARTIFACTS,
    filename="pipeline_artifacts_2021_08_04T15_46_02.tar",
)
# an output file
INGESTION_AF_ANCILLARY = AncillaryProduct(
    type=AncillaryProductType.INGESTION_ARTIFACTS,
    filename="ingestion_artifacts_2021_08_04T01_57_08.564.tar",
)
ANCILLARY_PRODUCTS = [WEBLOG_ANCILLARY, INGESTION_AF_ANCILLARY, PIPELINE_AF_ANCILLARY]

STAGING_DIR_FILES = [
    "aux_image_metadata.json",
    "oussid.J1522+3934_sci.K_band.cont.I.pbcor.tt0.fits",
    "oussid.J1522+3934_sci.X_band.cont.I.pbcor.tt0.fits",
]


def populate_fake_tmpx_ratuqh_ingest_path(staging_source_dir: Path, is_final: bool = False) -> List[Path]:
    """
    Make a bunch of fake files that should result in the example manifest.
    If this is version 2 or later of a standard calibration, include the initial weblog.

    :return:
    """
    fake_files_to_create = [ADDITIONAL_METADATA_FILENAME]

    for product in OUTPUT_GROUP_SCIENCE_PRODUCT_ANCILLARIES_K:
        fake_files_to_create.append(product.filename)

    for product in OUTPUT_GROUP_SCIENCE_PRODUCT_ANCILLARIES_X:
        fake_files_to_create.append(product.filename)

    fake_files_to_create.append(PIPELINE_AF_ANCILLARY.filename)
    fake_files_to_create.append(WEBLOG_ANCILLARY.filename)
    if is_final:
        fake_files_to_create.append(INIT_WEBLOG_ANCILLARY.filename)

    fake_files_to_create.append(OUTPUT_SCIENCE_PRODUCT_K.filename)
    fake_files_to_create.append(OUTPUT_SCIENCE_PRODUCT_X.filename)

    files = []
    for filename in fake_files_to_create:
        file = staging_source_dir / filename
        file.touch()
        files.append(file)

    assert len(files) == len(fake_files_to_create)

    return files


# -----------------------------
# SDM manifest test data
# -----------------------------
SDM_FILE_LIST = [
    "Antenna.xml",
    "ASDM.xml",
    "CalData.xml",
    "CalDevice.xml",
    "CalPointing.xml",
    "CalReduction.xml",
    "ConfigDescription.xml",
    "CorrelatorMode.xml",
    "DataDescription.xml",
    "Doppler.xml",
    "Ephemeris.xml",
    "ExecBlock.xml",
    "Feed.xml",
    "Field.xml",
    "Flag.xml",
    "Main.xml",
    "Pointing.xml",
    "PointingModel.xml",
    "Polarization.xml",
    "Processor.xml",
    "Receiver.xml",
    "SBSummary.xml",
    "Scan.xml",
    "Source.xml",
    "SpectralWindow.xml",
    "State.xml",
    "Station.xml",
    "Subscan.xml",
    "SwitchCycle.xml",
    "SysCal.xml",
    "SysPower.bin",
    "Weather.xml",
]
# From file ./examples/full_curation_evla_eb_manifest.json
EVLA_EB_NAME = "19A-001.sb1234567890.eb233423545632.54321.894327984569"
EVLA_EB_LOCATOR = "uid://I/am/a/locator"


def populate_fake_evla_eb_curator_source_path(staging_dir: Path) -> list[Path]:
    eb_dir = staging_dir / EVLA_EB_NAME
    eb_dir.mkdir()
    for sdm_filename in SDM_FILE_LIST:
        sdm_file_path = eb_dir / sdm_filename
        sdm_file_path.touch()
    # Only really care about the directory for the manifest
    return [eb_dir]


def find_ingestion_artifacts_tar(staging_source_dir: Path):
    """
    There should be an ingestion artifacts tar after manifest creation.

    :param staging_source_dir:
    :return:
    """
    ing_artifacts_tars = [
        file
        for file in staging_source_dir.iterdir()
        if file.name.startswith(AncillaryProductType.INGESTION_ARTIFACTS.value) and file.name.endswith(TARFILE_EXT)
    ]
    assert len(ing_artifacts_tars) == 1
    return ing_artifacts_tars[0]