test_endpoints.py

import json
import os
import re
import subprocess
import sys

import bson
import pytest
import requests
from dagster import build_op_context
from starlette import status
from tenacity import wait_random_exponential, retry
from toolz import get_in

from nmdc_runtime.api.core.auth import get_password_hash
from nmdc_runtime.api.core.metadata import df_from_sheet_in, _validate_changesheet
from nmdc_runtime.api.core.util import generate_secret, dotted_path_for
from nmdc_runtime.api.db.mongo import get_mongo_db, mongorestore_from_dir
from nmdc_runtime.api.endpoints.util import persist_content_and_get_drs_object
from nmdc_runtime.api.models.job import Job, JobOperationMetadata
from nmdc_runtime.api.models.metadata import ChangesheetIn
from nmdc_runtime.api.models.site import SiteInDB, SiteClientInDB
from nmdc_runtime.api.models.user import UserInDB, UserIn, User
from nmdc_runtime.site.ops import materialize_alldocs
from nmdc_runtime.site.repository import run_config_frozen__normal_env
from nmdc_runtime.site.resources import get_mongo, RuntimeApiSiteClient, mongo_resource
from nmdc_runtime.util import REPO_ROOT_DIR, ensure_unique_id_indexes
from tests.test_util import download_and_extract_tar
from tests.test_ops.test_ops import op_context as test_op_context

TEST_MONGODUMPS_DIR = REPO_ROOT_DIR.joinpath("tests", "nmdcdb")
SCHEMA_COLLECTIONS_MONGODUMP_ARCHIVE_BASENAME = (
    "nmdc-prod-schema-collections__2024-07-29_20-12-07"
)
SCHEMA_COLLECTIONS_MONGODUMP_ARCHIVE_URL = (
    "https://portal.nersc.gov/cfs/m3408/meta/mongodumps/"
    f"{SCHEMA_COLLECTIONS_MONGODUMP_ARCHIVE_BASENAME}.tar"
)  # 84MB. Should be < 100MB.


def ensure_local_mongodump_exists():
    dump_dir = TEST_MONGODUMPS_DIR.joinpath(
        SCHEMA_COLLECTIONS_MONGODUMP_ARCHIVE_BASENAME
    )
    if not os.path.exists(dump_dir):
        download_and_extract_tar(
            url=SCHEMA_COLLECTIONS_MONGODUMP_ARCHIVE_URL, extract_to=TEST_MONGODUMPS_DIR
        )
    else:
        print(f"local mongodump already exists at {TEST_MONGODUMPS_DIR}")
    return dump_dir


def ensure_schema_collections_and_alldocs():
    # Return if `alldocs` collection has already been materialized.
    mdb = get_mongo_db()
    if mdb.alldocs.estimated_document_count() > 0:
        print(
            "ensure_schema_collections_and_alldocs: `alldocs` collection already materialized"
        )
        return

    dump_dir = ensure_local_mongodump_exists()
    mongorestore_from_dir(mdb, dump_dir, skip_collections=["functional_annotation_agg"])
    ensure_unique_id_indexes(mdb)
    print("materializing alldocs...")
    materialize_alldocs(
        build_op_context(
            resources={
                "mongo": mongo_resource.configured(
                    {
                        "dbname": os.getenv("MONGO_DBNAME"),
                        "host": os.getenv("MONGO_HOST"),
                        "password": os.getenv("MONGO_PASSWORD"),
                        "username": os.getenv("MONGO_USERNAME"),
                    }
                )
            }
        )
    )


def ensure_test_resources(mdb):
    username = "testuser"
    password = generate_secret()
    site_id = "testsite"
    mdb.users.replace_one(
        {"username": username},
        UserInDB(
            username=username,
            hashed_password=get_password_hash(password),
            site_admin=[site_id],
        ).model_dump(exclude_unset=True),
        upsert=True,
    )

    client_id = "testsite-testclient"
    client_secret = generate_secret()
    mdb.sites.replace_one(
        {"id": site_id},
        SiteInDB(
            id=site_id,
            clients=[
                SiteClientInDB(
                    id=client_id,
                    hashed_secret=get_password_hash(client_secret),
                )
            ],
        ).model_dump(),
        upsert=True,
    )
    wf_id = "test"
    job_id = "nmdc:fk0jb83"
    prev_ops = {"metadata.job.id": job_id, "metadata.site_id": site_id}
    mdb.operations.delete_many(prev_ops)
    job = Job(**{"id": job_id, "workflow": {"id": wf_id}, "config": {}, "claims": []})
    mdb.jobs.replace_one(
        {"id": job_id}, job.model_dump(exclude_unset=True), upsert=True
    )
    mdb["minter.requesters"].replace_one({"id": site_id}, {"id": site_id}, upsert=True)
    ensure_schema_collections_and_alldocs()
    return {
        "site_client": {
            "site_id": site_id,
            "client_id": client_id,
            "client_secret": client_secret,
        },
        "user": {"username": username, "password": password},
        "job": job.model_dump(exclude_unset=True),
    }


@pytest.mark.skip(reason="Skipping because test causes suite to hang")
def test_update_operation():
    mdb = get_mongo(run_config_frozen__normal_env).db
    rs = ensure_test_resources(mdb)
    client = RuntimeApiSiteClient(base_url=os.getenv("API_HOST"), **rs["site_client"])
    rv = client.claim_job(rs["job"]["id"])
    assert "id" in rv.json()
    op = rv.json()
    new_op = client.update_operation(op["id"], {"metadata": {"foo": "bar"}}).json()
    assert get_in(["metadata", "site_id"], new_op) == rs["site_client"]["site_id"]
    assert get_in(["metadata", "job", "id"], new_op) == rs["job"]["id"]
    assert get_in(["metadata", "model"], new_op) == dotted_path_for(
        JobOperationMetadata
    )


@pytest.mark.skip(reason="Skipping because test causes suite to hang")
def test_create_user():
    mdb = get_mongo(run_config_frozen__normal_env).db
    rs = ensure_test_resources(mdb)
    base_url = os.getenv("API_HOST")

    @retry(wait=wait_random_exponential(multiplier=1, max=60))
    def get_token():
        """

        Randomly wait up to 2^x * 1 seconds between each retry until the range reaches 60
        seconds, then randomly up to 60 seconds afterwards

        """

        _rv = requests.post(
            base_url + "/token",
            data={
                "grant_type": "password",
                "username": rs["user"]["username"],
                "password": rs["user"]["password"],
            },
        )
        token_response = _rv.json()
        return token_response["access_token"]

    headers = {"Authorization": f"Bearer {get_token()}"}

    user_in = UserIn(username="foo", password=generate_secret())
    mdb.users.delete_one({"username": user_in.username})
    mdb.users.update_one(
        {"username": rs["user"]["username"]},
        {"$addToSet": {"site_admin": "nmdc-runtime-useradmin"}},
    )
    rv = requests.request(
        "POST",
        url=(base_url + "/users"),
        headers=headers,
        json=user_in.model_dump(exclude_unset=True),
    )

    try:
        assert rv.status_code == status.HTTP_201_CREATED
        User(**rv.json())
    finally:
        mdb.users.delete_one({"username": user_in.username})
        mdb.users.update_one(
            {"username": rs["user"]["username"]},
            {"$pull": {"site_admin": "nmdc-runtime-useradmin"}},
        )


@pytest.fixture
def api_site_client():
    mdb = get_mongo_db()
    rs = ensure_test_resources(mdb)
    return RuntimeApiSiteClient(base_url=os.getenv("API_HOST"), **rs["site_client"])


def test_metadata_validate_json_0(api_site_client):
    rv = api_site_client.request(
        "POST",
        "/metadata/json:validate",
        {
            "field_research_site_set": [
                {"id": "nmdc:frsite-11-s2dqk408", "name": "BESC-470-CL2_38_23"},
                {"id": "nmdc:frsite-11-s2dqk408", "name": "BESC-470-CL2_38_23"},
                {"id": "nmdc:frsite-11-s2dqk408", "name": "BESC-470-CL2_38_23"},
            ]
        },
    )
    assert rv.json()["result"] == "errors"


def test_metadata_validate_json_empty_collection(api_site_client):
    rv = api_site_client.request(
        "POST",
        "/metadata/json:validate",
        {"study_set": []},
    )
    assert rv.json()["result"] != "errors"


def test_metadata_validate_json_with_type_attribute(api_site_client):
    rv = api_site_client.request(
        "POST",
        "/metadata/json:validate",
        {"study_set": [], "@type": "Database"},
    )
    assert rv.json()["result"] != "errors"
    rv = api_site_client.request(
        "POST",
        "/metadata/json:validate",
        {"study_set": [], "@type": "nmdc:Database"},
    )
    assert rv.json()["result"] != "errors"


def test_metadata_validate_json_with_unknown_collection(api_site_client):
    rv = api_site_client.request(
        "POST",
        "/metadata/json:validate",
        {"studi_set": []},
    )
    assert rv.json()["result"] == "errors"


def test_submit_changesheet():
    sheet_in = ChangesheetIn(
        name="sheet",
        content_type="text/tab-separated-values",
        text="id\taction\tattribute\tvalue\nnmdc:bsm-12-7mysck21\tupdate\tassociated_studies\tnmdc:sty-11-pzmd0x14\n",
    )
    mdb = get_mongo_db()
    rs = ensure_test_resources(mdb)
    mdb.biosample_set.replace_one(
        {"id": "nmdc:bsm-12-7mysck21"},
        json.loads(
            (
                REPO_ROOT_DIR / "tests" / "files" / "nmdc_bsm-12-7mysck21.json"
            ).read_text()
        ),
        upsert=True,
    )
    mdb.study_set.replace_one(
        {"id": "nmdc:sty-11-pzmd0x14"},
        json.loads(
            (
                REPO_ROOT_DIR / "tests" / "files" / "nmdc_sty-11-pzmd0x14.json"
            ).read_text()
        ),
        upsert=True,
    )
    df_change = df_from_sheet_in(sheet_in, mdb)
    _ = _validate_changesheet(df_change, mdb)

    # clear objects collection to avoid strange duplicate key error.
    mdb.objects.delete_many({})

    drs_obj_doc = persist_content_and_get_drs_object(
        content=sheet_in.text,
        username=rs["user"]["username"],
        filename=re.sub(r"[^A-Za-z0-9._\-]", "_", sheet_in.name),
        content_type=sheet_in.content_type,
        description="changesheet",
        id_ns="changesheets",
    )
    mdb.objects.delete_one({"id": drs_obj_doc["id"]})
    assert True


def test_submit_workflow_executions(api_site_client):
    test_collection, test_id = "workflow_execution_set", "nmdc:wfmag-11-00jn7876.1"
    test_payload = {
        test_collection: [
            {
                "id": test_id,
                "name": "Metagenome Assembled Genomes Analysis Activity for nmdc:wfmag-11-00jn7876.1",
                "started_at_time": "2023-07-30T21:31:56.387227+00:00",
                "ended_at_time": "2023-07-30T21:34:32.750008+00:00",
                "was_informed_by": "nmdc:omprc-11-7yj0jg57",
                "execution_resource": "NERSC-Perlmutter",
                "git_url": "https://github.com/microbiomedata/metaMAGs",
                "has_input": [
                    "nmdc:dobj-11-yjp1xw52",
                    "nmdc:dobj-11-3av14y79",
                    "nmdc:dobj-11-wa5pnq42",
                    "nmdc:dobj-11-nexa9703",
                    "nmdc:dobj-11-j13n8739",
                    "nmdc:dobj-11-116fa706",
                    "nmdc:dobj-11-60d0na51",
                    "nmdc:dobj-11-2vbz7538",
                    "nmdc:dobj-11-1t48mn65",
                    "nmdc:dobj-11-1cvwk224",
                    "nmdc:dobj-11-cdna6f90",
                    "nmdc:dobj-11-4vb3ww76",
                    "nmdc:dobj-11-xv4qd072",
                    "nmdc:dobj-11-m7p3sb10",
                    "nmdc:dobj-11-j0t1rv33",
                ],
                "has_output": [
                    "nmdc:dobj-11-k5ad4209",
                    "nmdc:dobj-11-bw8nqt30",
                    "nmdc:dobj-11-199t2777",
                    "nmdc:dobj-11-2qfh8476",
                    "nmdc:dobj-11-fcsvq172",
                ],
                "type": "nmdc:MagsAnalysis",
                "version": "v1.0.6",
                "mags_list": [],
            }
        ]
    }
    mdb = get_mongo_db()
    if doc_to_restore := mdb[test_collection].find_one({"id": test_id}):
        mdb[test_collection].delete_one({"id": test_id})
    rv = api_site_client.request(
        "POST",
        "/workflows/workflow_executions",
        test_payload,
    )
    assert rv.json() == {"message": "jobs accepted"}

    # check that the document was added to the database
    # clean up database before `assert` to be sure that cleanup happens even if test fails.
    rv = api_site_client.request("GET", f"/nmdcschema/ids/{test_id}")
    mdb[test_collection].delete_one({"id": test_id})
    if doc_to_restore:
        mdb[test_collection].insert_one(doc_to_restore)
    assert "id" in rv.json() and rv.json()["id"] == test_id


def test_get_class_name_and_collection_names_by_doc_id():
    base_url = os.getenv("API_HOST")

    # Seed the database.
    mdb = get_mongo_db()
    study_set_collection = mdb.get_collection(name="study_set")
    fake_doc = dict(id="nmdc:sty-1-foobar")
    study_set_collection.replace_one(fake_doc, fake_doc, upsert=True)

    # Valid `id`, and the document exists in database.
    id_ = "nmdc:sty-1-foobar"
    response = requests.request(
        "GET", f"{base_url}/nmdcschema/ids/{id_}/collection-name"
    )
    body = response.json()
    assert response.status_code == 200
    assert body["id"] == id_
    assert body["collection_name"] == "study_set"

    # Valid `id`, but the document does not exist in database.
    id_ = "nmdc:sty-1-bazqux"
    response = requests.request(
        "GET", f"{base_url}/nmdcschema/ids/{id_}/collection-name"
    )
    assert response.status_code == 404

    # Invalid `id` (because "foo" is an invalid typecode).
    id_ = "nmdc:foo-1-foobar"
    response = requests.request(
        "GET", f"{base_url}/nmdcschema/ids/{id_}/collection-name"
    )
    assert response.status_code == 404


def test_find_data_objects_for_study(api_site_client):
    ensure_schema_collections_and_alldocs()
    rv = api_site_client.request(
        "GET",
        "/data_objects/study/nmdc:sty-11-hdd4bf83",
    )
    assert len(rv.json()) >= 60

    
def test_find_workflow_executions(api_site_client):
    test_collection, test_id = (
        "workflow_execution_set",
        "nmdc:wfmgan-11-ndgg7v31.1",
    )
    test_doc = {
        "id": test_id,
        "name": "Annotation Activity for nmdc:wfmgan-11-ndgg7v31.1",
        "started_at_time": "2023-03-07T22:54:55.914797+00:00",
        "ended_at_time": "2023-03-07T22:54:55.914832+00:00",
        "was_informed_by": "nmdc:omprc-11-m0dd0851",
        "execution_resource": "JGI",
        "git_url": "https://github.com/microbiomedata/mg_annotation",
        "has_input": ["nmdc:dobj-11-a2w9zz17"],
        "has_output": [
            "nmdc:dobj-11-ej6gdk68",
            "nmdc:dobj-11-x50rg190",
            "nmdc:dobj-11-dk50fw35",
            "nmdc:dobj-11-x1b6f376",
            "nmdc:dobj-11-rhw29h15",
            "nmdc:dobj-11-kpgxx958",
            "nmdc:dobj-11-zzp8vt52",
            "nmdc:dobj-11-zvahqs54",
            "nmdc:dobj-11-38xhmg17",
            "nmdc:dobj-11-n8vqe154",
            "nmdc:dobj-11-55021k46",
            "nmdc:dobj-11-xxfaj025",
            "nmdc:dobj-11-55gp7g13",
            "nmdc:dobj-11-wmr56107",
            "nmdc:dobj-11-esyjjz10",
            "nmdc:dobj-11-xpbzyc98",
            "nmdc:dobj-11-8zpavw69",
            "nmdc:dobj-11-72df2803",
            "nmdc:dobj-11-1j6f8010",
            "nmdc:dobj-11-tkm6xd10",
            "nmdc:dobj-11-khw2qa20",
            "nmdc:dobj-11-mmy40b21",
            "nmdc:dobj-11-kq3rt657",
            "nmdc:dobj-11-5yekv009",
        ],
        "type": "nmdc:MetagenomeAnnotation",
        "version": "v1.0.2-beta",
        "gold_analysis_project_identifiers": [],
    }

    mdb = get_mongo_db()
    mdb[test_collection].replace_one({"id": test_id}, test_doc, upsert=True)
    rv = api_site_client.request(
        "GET", "/workflow_executions", params_or_json_data={"filter": f"id:{test_id}"}
    )
    assert "meta" in rv.json() and rv.json()["meta"]["count"] == 1
    rv = api_site_client.request(
        "GET",
        f"/workflow_executions/{test_id}",
    )
    assert "id" in rv.json() and rv.json()["id"] == test_id