datahub/smoke-test/tests/containers/containers_test.py

import pytest
import time
from tests.utils import FRONTEND_ENDPOINT
from tests.utils import GMS_ENDPOINT
from tests.utils import ingest_file_via_rest
from tests.utils import delete_urns_from_file

@pytest.fixture(scope="module", autouse=False)
def ingest_cleanup_data(request):
    print("ingesting containers test data")
    ingest_file_via_rest("tests/containers/data.json")
    yield
    print("removing containers test data")
    delete_urns_from_file("tests/containers/data.json")

@pytest.mark.dependency()
def test_healthchecks(wait_for_healthchecks):
    # Call to wait_for_healthchecks fixture will do the actual functionality.
    pass

@pytest.mark.dependency(depends=["test_healthchecks"])
def test_get_full_container(frontend_session, ingest_cleanup_data):

    container_urn = "urn:li:container:SCHEMA"
    container_name = "datahub_schema"
    container_description = "The DataHub schema"
    editable_container_description = "custom description"
    dataset_urn = "urn:li:dataset:(urn:li:dataPlatform:hive,SampleHiveDataset,PROD)"

    # Get a full container
    get_container_json = {
        "query": """query container($urn: String!) {\n
            container(urn: $urn) {\n
              urn\n
              type\n
              platform {\n
                urn\n
                properties{\n
                  displayName\n
                }\n
              }\n
              container {\n
                urn\n
                properties {\n
                  name\n
                  description\n
                }\n
              }\n
              properties {\n
                name\n
                description\n
              }\n
              editableProperties {\n
                description\n
              }\n
              ownership {\n
                owners {\n
                  owner {\n
                    ...on CorpUser {\n
                      urn\n
                    }\n
                  }\n
                }\n
              }\n
              institutionalMemory {\n
                elements {\n
                  url\n
                }\n
              }\n
              tags {\n
                tags {\n
                  tag {\n
                    urn\n
                  }\n
                }\n
              }\n
              glossaryTerms {\n
                terms {\n
                  term {\n
                    urn\n
                  }\n
                }\n
              }\n
              subTypes {\n
                typeNames\n
              }\n
              entities(input: {}) {\n
                total\n
                searchResults {\n
                  entity {\n
                    ...on Dataset {\n
                      urn\n
                    }\n
                  }\n
                }\n
              }\n
            }\n
        }""",
        "variables": {
          "urn": container_urn
        }
    }

    response = frontend_session.post(
        f"{FRONTEND_ENDPOINT}/api/v2/graphql", json=get_container_json
    )
    response.raise_for_status()
    res_data = response.json()

    assert res_data
    assert res_data["data"]
    assert res_data["data"]["container"] is not None
    assert "errors" not in res_data

    container = res_data["data"]["container"]
    assert container["urn"] == container_urn
    assert container["type"] == "CONTAINER"
    assert container["platform"]["urn"] == "urn:li:dataPlatform:mysql"
    assert container["properties"]["name"] == container_name
    assert container["properties"]["description"] == container_description
    assert container["subTypes"]["typeNames"][0] == "Schema"
    assert container["editableProperties"]["description"] == editable_container_description
    assert container["ownership"] is None
    assert container["institutionalMemory"] is None
    assert container["tags"] is None
    assert container["glossaryTerms"] is None

@pytest.mark.dependency(depends=["test_healthchecks", "test_get_full_container"])
def test_get_parent_container(frontend_session, ingest_cleanup_data):

    dataset_urn = "urn:li:dataset:(urn:li:dataPlatform:hive,SampleHiveDataset,PROD)"

    # Get count of existing secrets
    get_dataset_json = {
        "query": """query dataset($urn: String!) {\n
          dataset(urn: $urn) {\n
            urn\n
            container {\n
              urn\n
              properties {\n
                name\n
              }\n
            }\n
          }\n
        }""",
        "variables": {
          "urn": dataset_urn
        }
    }

    response = frontend_session.post(
        f"{FRONTEND_ENDPOINT}/api/v2/graphql", json=get_dataset_json
    )
    response.raise_for_status()
    res_data = response.json()

    assert res_data
    assert res_data["data"]
    assert res_data["data"]["dataset"] is not None
    assert "errors" not in res_data

    dataset = res_data["data"]["dataset"]
    assert dataset["container"]["properties"]["name"] == "datahub_schema"

@pytest.mark.dependency(depends=["test_healthchecks", "test_get_full_container"])
def test_update_container(frontend_session, ingest_cleanup_data):

    container_urn = "urn:li:container:SCHEMA"

    new_tag = "urn:li:tag:Test"

    add_tag_json = {
        "query": """mutation addTag($input: TagAssociationInput!) {\n
            addTag(input: $input)
        }""",
        "variables": {
            "input": {
              "tagUrn": new_tag,
              "resourceUrn": container_urn,
            }
        }
    }

    response = frontend_session.post(
        f"{FRONTEND_ENDPOINT}/api/v2/graphql", json=add_tag_json
    )
    response.raise_for_status()
    res_data = response.json()

    assert res_data
    assert res_data["data"]
    assert res_data["data"]["addTag"] is True

    new_term = "urn:li:glossaryTerm:Term"

    add_term_json = {
        "query": """mutation addTerm($input: TermAssociationInput!) {\n
            addTerm(input: $input)
        }""",
        "variables": {
            "input": {
              "termUrn": new_term,
              "resourceUrn": container_urn,
            }
        }
    }

    response = frontend_session.post(
        f"{FRONTEND_ENDPOINT}/api/v2/graphql", json=add_term_json
    )
    response.raise_for_status()
    res_data = response.json()

    assert res_data
    assert res_data["data"]
    assert res_data["data"]["addTerm"] is True

    new_owner = "urn:li:corpuser:jdoe"

    add_owner_json = {
        "query": """mutation addOwner($input: AddOwnerInput!) {\n
            addOwner(input: $input)
        }""",
        "variables": {
            "input": {
              "ownerUrn": new_owner,
              "resourceUrn": container_urn,
              "ownerEntityType": "CORP_USER"
            }
        }
    }

    response = frontend_session.post(
        f"{FRONTEND_ENDPOINT}/api/v2/graphql", json=add_owner_json
    )
    response.raise_for_status()
    res_data = response.json()

    assert res_data
    assert res_data["data"]
    assert res_data["data"]["addOwner"] is True

    new_link = "https://www.test.com"

    add_link_json = {
        "query": """mutation addLink($input: AddLinkInput!) {\n
            addLink(input: $input)
        }""",
        "variables": {
            "input": {
              "linkUrl": new_link,
              "resourceUrn": container_urn,
              "label": "Label"
            }
        }
    }

    response = frontend_session.post(
        f"{FRONTEND_ENDPOINT}/api/v2/graphql", json=add_link_json
    )
    response.raise_for_status()
    res_data = response.json()

    assert res_data
    assert res_data["data"]
    assert res_data["data"]["addLink"] is True

    new_description = "New description"

    update_description_json = {
        "query": """mutation updateDescription($input: DescriptionUpdateInput!) {\n
            updateDescription(input: $input)
        }""",
        "variables": {
            "input": {
              "description": new_description,
              "resourceUrn": container_urn,
            }
        }
    }

    response = frontend_session.post(
        f"{FRONTEND_ENDPOINT}/api/v2/graphql", json=update_description_json
    )
    response.raise_for_status()
    res_data = response.json()

    assert res_data
    assert res_data["data"]
    assert res_data["data"]["updateDescription"] is True

    # Now fetch the container to ensure it was updated
   # Get the container
    get_container_json = {
        "query": """query container($urn: String!) {\n
           container(urn: $urn) {\n
              editableProperties {\n
                description\n
              }\n
              ownership {\n
                owners {\n
                  owner {\n
                    ...on CorpUser {\n
                      urn\n
                    }\n
                  }\n
                }\n
              }\n
              institutionalMemory {\n
                elements {\n
                  url\n
                }\n
              }\n
              tags {\n
                tags {\n
                  tag {\n
                    urn\n
                  }\n
                }\n
              }\n
              glossaryTerms {\n
                terms {\n
                  term {\n
                    urn\n
                  }\n
                }\n
              }\n
            }\n
        }""",
        "variables": {
          "urn": container_urn
        }
    }

    response = frontend_session.post(
        f"{FRONTEND_ENDPOINT}/api/v2/graphql", json=get_container_json
    )
    response.raise_for_status()
    res_data = response.json()

    assert res_data
    assert res_data["data"]
    assert res_data["data"]["container"] is not None
    assert "errors" not in res_data

    container = res_data["data"]["container"]
    assert container["editableProperties"]["description"] == new_description
    assert container["ownership"]["owners"][0]["owner"]["urn"] == new_owner
    assert container["institutionalMemory"]["elements"][0]["url"] == new_link
    assert container["tags"]["tags"][0]["tag"]["urn"] == new_tag
    assert container["glossaryTerms"]["terms"][0]["term"]["urn"] == new_term