fix(test): add cleanup in tests, make urls configurable (#5287)

smoke-test/test_rapid.py

@@ -1,25 +1,22 @@
 import time
-import urllib
-from typing import Any, Dict, Optional, cast
 
 import pytest
 import requests
-
 from datahub.cli.docker import check_local_docker_containers
-from datahub.ingestion.run.pipeline import Pipeline
-from datahub.ingestion.source.state.checkpoint import Checkpoint
-from tests.utils import ingest_file_via_rest
+
+from tests.utils import get_frontend_url, ingest_file_via_rest
 
 bootstrap_small = "test_resources/bootstrap_single.json"
 bootstrap_small_2 = "test_resources/bootstrap_single2.json"
-FRONTEND_ENDPOINT = "http://localhost:9002"
+
 
 @pytest.fixture(scope="session")
 def wait_for_healthchecks():
     # Simply assert that everything is healthy, but don't wait.
     assert not check_local_docker_containers()
     yield
 
+
 @pytest.fixture(scope="session")
 def frontend_session(wait_for_healthchecks):
     session = requests.Session()
@@ -28,19 +25,18 @@ def frontend_session(wait_for_healthchecks):
         "Content-Type": "application/json",
     }
     data = '{"username":"datahub", "password":"datahub"}'
-    response = session.post(
-        f"{FRONTEND_ENDPOINT}/logIn", headers=headers, data=data
-    )
+    response = session.post(f"{get_frontend_url()}/logIn", headers=headers, data=data)
     response.raise_for_status()
 
     yield session
 
+
 def test_ingestion_via_rest_rapid(frontend_session, wait_for_healthchecks):
     ingest_file_via_rest(bootstrap_small)
     ingest_file_via_rest(bootstrap_small_2)
-    urn = f"urn:li:dataset:(urn:li:dataPlatform:testPlatform,testDataset,PROD)"
+    urn = "urn:li:dataset:(urn:li:dataPlatform:testPlatform,testDataset,PROD)"
     json = {
-            "query": """query getDataset($urn: String!) {\n
+        "query": """query getDataset($urn: String!) {\n
                 dataset(urn: $urn) {\n
                     urn\n
                     name\n
@@ -70,15 +66,11 @@ def test_ingestion_via_rest_rapid(frontend_session, wait_for_healthchecks):
                     }\n
                 }\n
             }""",
-            "variables": {
-                "urn": urn
-            }
-        }
+        "variables": {"urn": urn},
+    }
     #
     time.sleep(2)
-    response = frontend_session.post(
-        f"{FRONTEND_ENDPOINT}/api/v2/graphql", json=json
-    )
+    response = frontend_session.post(f"{get_frontend_url()}/api/v2/graphql", json=json)
     response.raise_for_status()
     res_data = response.json()
 

smoke-test/tests/assertions/assertions_test.py

@@ -4,7 +4,6 @@
 
 import pytest
 import requests
-
 from datahub.cli.docker import check_local_docker_containers
 from datahub.emitter.mce_builder import make_dataset_urn, make_schema_field_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
@@ -25,11 +24,7 @@
     PartitionSpecClass,
     PartitionTypeClass,
 )
-from tests.utils import ingest_file_via_rest
-from tests.utils import delete_urns_from_file
-
-
-GMS_ENDPOINT = "http://localhost:8080"
+from tests.utils import delete_urns_from_file, get_gms_url, ingest_file_via_rest
 
 restli_default_headers = {
     "X-RestLi-Protocol-Version": "2.0.0",
@@ -253,6 +248,7 @@ def test_healthchecks(wait_for_healthchecks):
 def test_run_ingestion(generate_test_data):
     ingest_file_via_rest(generate_test_data)
 
+
 @pytest.mark.dependency(depends=["test_healthchecks", "test_run_ingestion"])
 def test_gms_get_latest_assertions_results_by_partition():
     urn = make_dataset_urn("postgres", "foo")
@@ -294,7 +290,7 @@ def test_gms_get_latest_assertions_results_by_partition():
         }
     )
     response = requests.post(
-        f"{GMS_ENDPOINT}/analytics?action=getTimeseriesStats",
+        f"{get_gms_url()}/analytics?action=getTimeseriesStats",
         data=query,
         headers=restli_default_headers,
     )
@@ -325,7 +321,7 @@ def test_gms_get_assertions_on_dataset():
     """lists all assertion urns including those which may not have executed"""
     urn = make_dataset_urn("postgres", "foo")
     response = requests.get(
-        f"{GMS_ENDPOINT}/relationships?direction=INCOMING&urn={urllib.parse.quote(urn)}&types=Asserts"
+        f"{get_gms_url()}/relationships?direction=INCOMING&urn={urllib.parse.quote(urn)}&types=Asserts"
     )
 
     response.raise_for_status()
@@ -339,7 +335,7 @@ def test_gms_get_assertions_on_dataset_field():
     dataset_urn = make_dataset_urn("postgres", "foo")
     field_urn = make_schema_field_urn(dataset_urn, "col1")
     response = requests.get(
-        f"{GMS_ENDPOINT}/relationships?direction=INCOMING&urn={urllib.parse.quote(field_urn)}&types=Asserts"
+        f"{get_gms_url()}/relationships?direction=INCOMING&urn={urllib.parse.quote(field_urn)}&types=Asserts"
     )
 
     response.raise_for_status()
@@ -351,7 +347,7 @@ def test_gms_get_assertions_on_dataset_field():
 def test_gms_get_assertion_info():
     assertion_urn = "urn:li:assertion:2d3b06a6e77e1f24adc9860a05ea089b"
     response = requests.get(
-        f"{GMS_ENDPOINT}/aspects/{urllib.parse.quote(assertion_urn)}\
+        f"{get_gms_url()}/aspects/{urllib.parse.quote(assertion_urn)}\
             ?aspect=assertionInfo&version=0",
         headers=restli_default_headers,
     )
@@ -364,4 +360,4 @@ def test_gms_get_assertion_info():
     assert data["aspect"]["com.linkedin.assertion.AssertionInfo"]["type"] == "DATASET"
     assert data["aspect"]["com.linkedin.assertion.AssertionInfo"]["datasetAssertion"][
         "scope"
-    ]
+    ]

smoke-test/tests/cli/datahub_graph_test.py

@@ -1,14 +1,7 @@
-import time
-
 import pytest
 from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph
-from datahub.metadata.schema_classes import SchemaMetadataClass, KafkaSchemaClass
-from tests.utils import (
-    FRONTEND_ENDPOINT,
-    GMS_ENDPOINT,
-    delete_urns_from_file,
-    ingest_file_via_rest,
-)
+from datahub.metadata.schema_classes import KafkaSchemaClass, SchemaMetadataClass
+from tests.utils import delete_urns_from_file, ingest_file_via_rest
 
 
 @pytest.fixture(scope="module", autouse=False)
@@ -38,5 +31,7 @@ def test_get_aspect_v2(frontend_session, ingest_cleanup_data):
     assert schema_metadata.platform == "urn:li:dataPlatform:kafka"
     assert isinstance(schema_metadata.platformSchema, KafkaSchemaClass)
     k_schema: KafkaSchemaClass = schema_metadata.platformSchema
-    assert k_schema.documentSchema == "{\"type\":\"record\",\"name\":\"SampleKafkaSchema\",\"namespace\":\"com.linkedin.dataset\",\"doc\":\"Sample Kafka dataset\",\"fields\":[{\"name\":\"field_foo\",\"type\":[\"string\"]},{\"name\":\"field_bar\",\"type\":[\"boolean\"]}]}"
-
+    assert (
+        k_schema.documentSchema
+        == '{"type":"record","name":"SampleKafkaSchema","namespace":"com.linkedin.dataset","doc":"Sample Kafka dataset","fields":[{"name":"field_foo","type":["string"]},{"name":"field_bar","type":["boolean"]}]}'
+    )

smoke-test/tests/conftest.py

@@ -1,22 +1,22 @@
 import os
-import time
 
 import pytest
 import requests
-import urllib
 from datahub.cli.docker import check_local_docker_containers
-from datahub.ingestion.run.pipeline import Pipeline
-from tests.utils import FRONTEND_ENDPOINT
+
+from tests.utils import get_frontend_url
 
 # Disable telemetry
 os.putenv("DATAHUB_TELEMETRY_ENABLED", "false")
 
+
 @pytest.fixture(scope="session")
 def wait_for_healthchecks():
     # Simply assert that everything is healthy, but don't wait.
     assert not check_local_docker_containers()
     yield
 
+
 @pytest.fixture(scope="session")
 def frontend_session(wait_for_healthchecks):
     session = requests.Session()
@@ -25,15 +25,14 @@ def frontend_session(wait_for_healthchecks):
         "Content-Type": "application/json",
     }
     data = '{"username":"datahub", "password":"datahub"}'
-    response = session.post(
-        f"{FRONTEND_ENDPOINT}/logIn", headers=headers, data=data
-    )
+    response = session.post(f"{get_frontend_url()}/logIn", headers=headers, data=data)
     response.raise_for_status()
 
     yield session
 
+
 # TODO: Determine whether we need this or not.
 @pytest.mark.dependency()
 def test_healthchecks(wait_for_healthchecks):
     # Call to wait_for_healthchecks fixture will do the actual functionality.
-    pass
+    pass