fix(lint): run black, isort

smoke-test/tests/assertions/assertions_test.py

@@ -2,28 +2,29 @@
 import urllib
 
 import pytest
-import requests_wrapper as requests
 import tenacity
 from datahub.emitter.mce_builder import make_dataset_urn, make_schema_field_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext, RecordEnvelope
 from datahub.ingestion.api.sink import NoopWriteCallback
 from datahub.ingestion.sink.file import FileSink, FileSinkConfig
-from datahub.metadata.com.linkedin.pegasus2avro.assertion import AssertionStdAggregation
-from datahub.metadata.schema_classes import (
-    AssertionInfoClass,
-    AssertionResultClass,
-    AssertionResultTypeClass,
-    AssertionRunEventClass,
-    AssertionRunStatusClass,
-    AssertionStdOperatorClass,
-    AssertionTypeClass,
-    DatasetAssertionInfoClass,
-    DatasetAssertionScopeClass,
-    PartitionSpecClass,
-    PartitionTypeClass,
-)
-from tests.utils import delete_urns_from_file, get_gms_url, ingest_file_via_rest, wait_for_healthcheck_util, get_sleep_info
+from datahub.metadata.com.linkedin.pegasus2avro.assertion import \
+    AssertionStdAggregation
+from datahub.metadata.schema_classes import (AssertionInfoClass,
+                                             AssertionResultClass,
+                                             AssertionResultTypeClass,
+                                             AssertionRunEventClass,
+                                             AssertionRunStatusClass,
+                                             AssertionStdOperatorClass,
+                                             AssertionTypeClass,
+                                             DatasetAssertionInfoClass,
+                                             DatasetAssertionScopeClass,
+                                             PartitionSpecClass,
+                                             PartitionTypeClass)
+
+import requests_wrapper as requests
+from tests.utils import (delete_urns_from_file, get_gms_url, get_sleep_info,
+                         ingest_file_via_rest, wait_for_healthcheck_util)
 
 restli_default_headers = {
     "X-RestLi-Protocol-Version": "2.0.0",

smoke-test/tests/browse/browse_test.py

@@ -1,9 +1,10 @@
 import time
 
 import pytest
-import requests_wrapper as requests
-from tests.utils import delete_urns_from_file, get_frontend_url, ingest_file_via_rest
 
+import requests_wrapper as requests
+from tests.utils import (delete_urns_from_file, get_frontend_url,
+                         ingest_file_via_rest)
 
 TEST_DATASET_1_URN = "urn:li:dataset:(urn:li:dataPlatform:kafka,test-browse-1,PROD)"
 TEST_DATASET_2_URN = "urn:li:dataset:(urn:li:dataPlatform:kafka,test-browse-2,PROD)"
@@ -51,7 +52,9 @@ def test_get_browse_paths(frontend_session, ingest_cleanup_data):
     # /prod -- There should be one entity
     get_browse_paths_json = {
         "query": get_browse_paths_query,
-        "variables": {"input": { "type": "DATASET", "path": ["prod"], "start": 0, "count": 100 } },
+        "variables": {
+            "input": {"type": "DATASET", "path": ["prod"], "start": 0, "count": 100}
+        },
     }
 
     response = frontend_session.post(
@@ -67,12 +70,19 @@ def test_get_browse_paths(frontend_session, ingest_cleanup_data):
 
     browse = res_data["data"]["browse"]
     print(browse)
-    assert browse["entities"] == [{ "urn": TEST_DATASET_3_URN }]
+    assert browse["entities"] == [{"urn": TEST_DATASET_3_URN}]
 
     # /prod/kafka1
     get_browse_paths_json = {
         "query": get_browse_paths_query,
-        "variables": {"input": { "type": "DATASET", "path": ["prod", "kafka1"], "start": 0, "count": 10 } },
+        "variables": {
+            "input": {
+                "type": "DATASET",
+                "path": ["prod", "kafka1"],
+                "start": 0,
+                "count": 10,
+            }
+        },
     }
 
     response = frontend_session.post(
@@ -88,16 +98,27 @@ def test_get_browse_paths(frontend_session, ingest_cleanup_data):
 
     browse = res_data["data"]["browse"]
     assert browse == {
-      "total": 3,
-      "entities": [{ "urn": TEST_DATASET_1_URN }, { "urn": TEST_DATASET_2_URN }, { "urn": TEST_DATASET_3_URN }],
-      "groups": [],
-      "metadata": { "path": ["prod", "kafka1"], "totalNumEntities": 0 }
+        "total": 3,
+        "entities": [
+            {"urn": TEST_DATASET_1_URN},
+            {"urn": TEST_DATASET_2_URN},
+            {"urn": TEST_DATASET_3_URN},
+        ],
+        "groups": [],
+        "metadata": {"path": ["prod", "kafka1"], "totalNumEntities": 0},
     }
 
     # /prod/kafka2
     get_browse_paths_json = {
         "query": get_browse_paths_query,
-        "variables": {"input": { "type": "DATASET", "path": ["prod", "kafka2"], "start": 0, "count": 10 } },
+        "variables": {
+            "input": {
+                "type": "DATASET",
+                "path": ["prod", "kafka2"],
+                "start": 0,
+                "count": 10,
+            }
+        },
     }
 
     response = frontend_session.post(
@@ -113,10 +134,8 @@ def test_get_browse_paths(frontend_session, ingest_cleanup_data):
 
     browse = res_data["data"]["browse"]
     assert browse == {
-      "total": 2,
-      "entities": [{ "urn": TEST_DATASET_1_URN }, { "urn": TEST_DATASET_2_URN }],
-      "groups": [],
-      "metadata": { "path": ["prod", "kafka2"], "totalNumEntities": 0 }
+        "total": 2,
+        "entities": [{"urn": TEST_DATASET_1_URN}, {"urn": TEST_DATASET_2_URN}],
+        "groups": [],
+        "metadata": {"path": ["prod", "kafka2"], "totalNumEntities": 0},
     }
-
-

smoke-test/tests/cli/datahub-cli.py

@@ -1,8 +1,11 @@
 import json
-import pytest
 from time import sleep
-from datahub.cli.cli_utils import guess_entity_type, post_entity, get_aspects_for_entity
+
+import pytest
+from datahub.cli.cli_utils import (get_aspects_for_entity, guess_entity_type,
+                                   post_entity)
 from datahub.cli.ingest_cli import get_session_and_host, rollback
+
 from tests.utils import ingest_file_via_rest, wait_for_writes_to_sync
 
 ingested_dataset_run_id = ""
@@ -24,58 +27,93 @@ def test_setup():
 
     session, gms_host = get_session_and_host()
 
-    assert "browsePaths" not in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["browsePaths"], typed=False)
-    assert "editableDatasetProperties" not in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False)
+    assert "browsePaths" not in get_aspects_for_entity(
+        entity_urn=dataset_urn, aspects=["browsePaths"], typed=False
+    )
+    assert "editableDatasetProperties" not in get_aspects_for_entity(
+        entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False
+    )
 
-    ingested_dataset_run_id = ingest_file_via_rest("tests/cli/cli_test_data.json").config.run_id
+    ingested_dataset_run_id = ingest_file_via_rest(
+        "tests/cli/cli_test_data.json"
+    ).config.run_id
     print("Setup ingestion id: " + ingested_dataset_run_id)
 
-    assert "browsePaths" in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["browsePaths"], typed=False)
+    assert "browsePaths" in get_aspects_for_entity(
+        entity_urn=dataset_urn, aspects=["browsePaths"], typed=False
+    )
 
     yield
 
     # Clean up
     rollback_url = f"{gms_host}/runs?action=rollback"
 
-    session.post(rollback_url, data=json.dumps({"runId": ingested_editable_run_id, "dryRun": False, "hardDelete": True}))
-    session.post(rollback_url, data=json.dumps({"runId": ingested_dataset_run_id, "dryRun": False, "hardDelete": True}))
+    session.post(
+        rollback_url,
+        data=json.dumps(
+            {"runId": ingested_editable_run_id, "dryRun": False, "hardDelete": True}
+        ),
+    )
+    session.post(
+        rollback_url,
+        data=json.dumps(
+            {"runId": ingested_dataset_run_id, "dryRun": False, "hardDelete": True}
+        ),
+    )
 
-    assert "browsePaths" not in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["browsePaths"], typed=False)
-    assert "editableDatasetProperties" not in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False)
+    assert "browsePaths" not in get_aspects_for_entity(
+        entity_urn=dataset_urn, aspects=["browsePaths"], typed=False
+    )
+    assert "editableDatasetProperties" not in get_aspects_for_entity(
+        entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False
+    )
 
 
 @pytest.mark.dependency()
 def test_rollback_editable():
     global ingested_dataset_run_id
     global ingested_editable_run_id
     platform = "urn:li:dataPlatform:kafka"
-    dataset_name = (
-        "test-rollback"
-    )
+    dataset_name = "test-rollback"
     env = "PROD"
     dataset_urn = f"urn:li:dataset:({platform},{dataset_name},{env})"
 
     session, gms_host = get_session_and_host()
 
     print("Ingested dataset id:", ingested_dataset_run_id)
     # Assert that second data ingestion worked
-    assert "browsePaths" in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["browsePaths"], typed=False)
+    assert "browsePaths" in get_aspects_for_entity(
+        entity_urn=dataset_urn, aspects=["browsePaths"], typed=False
+    )
 
     # Make editable change
-    ingested_editable_run_id = ingest_file_via_rest("tests/cli/cli_editable_test_data.json").config.run_id
+    ingested_editable_run_id = ingest_file_via_rest(
+        "tests/cli/cli_editable_test_data.json"
+    ).config.run_id
     print("ingested editable id:", ingested_editable_run_id)
     # Assert that second data ingestion worked
-    assert "editableDatasetProperties" in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False)
+    assert "editableDatasetProperties" in get_aspects_for_entity(
+        entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False
+    )
 
     # rollback ingestion 1
     rollback_url = f"{gms_host}/runs?action=rollback"
 
-    session.post(rollback_url, data=json.dumps({"runId": ingested_dataset_run_id, "dryRun": False, "hardDelete": False}))
+    session.post(
+        rollback_url,
+        data=json.dumps(
+            {"runId": ingested_dataset_run_id, "dryRun": False, "hardDelete": False}
+        ),
+    )
 
     # Allow async MCP processor to handle ingestions & rollbacks
     wait_for_writes_to_sync()
 
     # EditableDatasetProperties should still be part of the entity that was soft deleted.
-    assert "editableDatasetProperties" in get_aspects_for_entity(entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False)
+    assert "editableDatasetProperties" in get_aspects_for_entity(
+        entity_urn=dataset_urn, aspects=["editableDatasetProperties"], typed=False
+    )
     # But first ingestion aspects should not be present
-    assert "browsePaths" not in get_aspects_for_entity(entity_urn=dataset_urn, typed=False)
+    assert "browsePaths" not in get_aspects_for_entity(
+        entity_urn=dataset_urn, typed=False
+    )

smoke-test/tests/cli/datahub_graph_test.py

@@ -1,13 +1,11 @@
 import pytest
 import tenacity
 from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph
-from datahub.metadata.schema_classes import KafkaSchemaClass, SchemaMetadataClass
-from tests.utils import (
-    delete_urns_from_file,
-    get_gms_url,
-    get_sleep_info,
-    ingest_file_via_rest,
-)
+from datahub.metadata.schema_classes import (KafkaSchemaClass,
+                                             SchemaMetadataClass)
+
+from tests.utils import (delete_urns_from_file, get_gms_url, get_sleep_info,
+                         ingest_file_via_rest)
 
 sleep_sec, sleep_times = get_sleep_info()
 

smoke-test/tests/cli/delete_cmd/test_timeseries_delete.py

@@ -1,21 +1,22 @@
 import json
 import logging
+import sys
 import tempfile
 import time
-import sys
 from json import JSONDecodeError
 from typing import Any, Dict, List, Optional
 
-from click.testing import CliRunner, Result
-
 import datahub.emitter.mce_builder as builder
+from click.testing import CliRunner, Result
 from datahub.emitter.serialization_helper import pre_json_transform
 from datahub.entrypoints import datahub
 from datahub.metadata.schema_classes import DatasetProfileClass
+
+import requests_wrapper as requests
 from tests.aspect_generators.timeseries.dataset_profile_gen import \
     gen_dataset_profiles
-from tests.utils import get_strftime_from_timestamp_millis, wait_for_writes_to_sync
-import requests_wrapper as requests
+from tests.utils import (get_strftime_from_timestamp_millis,
+                         wait_for_writes_to_sync)
 
 logger = logging.getLogger(__name__)
 
@@ -33,6 +34,7 @@
 def sync_elastic() -> None:
     wait_for_writes_to_sync()
 
+
 def datahub_put_profile(dataset_profile: DatasetProfileClass) -> None:
     with tempfile.NamedTemporaryFile("w+t", suffix=".json") as aspect_file:
         aspect_text: str = json.dumps(pre_json_transform(dataset_profile.to_obj()))

smoke-test/tests/cli/ingest_cmd/test_timeseries_rollback.py

@@ -2,14 +2,14 @@
 import time
 from typing import Any, Dict, List, Optional
 
-from click.testing import CliRunner, Result
-
 import datahub.emitter.mce_builder as builder
+from click.testing import CliRunner, Result
 from datahub.emitter.serialization_helper import post_json_transform
 from datahub.entrypoints import datahub
 from datahub.metadata.schema_classes import DatasetProfileClass
-from tests.utils import ingest_file_via_rest, wait_for_writes_to_sync
+
 import requests_wrapper as requests
+from tests.utils import ingest_file_via_rest, wait_for_writes_to_sync
 
 runner = CliRunner(mix_stderr=False)
 

smoke-test/tests/cli/user_groups_cmd/test_group_cmd.py

@@ -1,14 +1,15 @@
 import json
 import sys
 import tempfile
+import time
 from typing import Any, Dict, Iterable, List
 
 import yaml
 from click.testing import CliRunner, Result
 from datahub.api.entities.corpgroup.corpgroup import CorpGroup
 from datahub.entrypoints import datahub
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
-import time
+
 import requests_wrapper as requests
 from tests.utils import wait_for_writes_to_sync
 

smoke-test/tests/conftest.py

@@ -2,8 +2,8 @@
 
 import pytest
 
-from tests.utils import wait_for_healthcheck_util, get_frontend_session
 from tests.test_result_msg import send_message
+from tests.utils import get_frontend_session, wait_for_healthcheck_util
 
 # Disable telemetry
 os.environ["DATAHUB_TELEMETRY_ENABLED"] = "false"
@@ -28,5 +28,5 @@ def test_healthchecks(wait_for_healthchecks):
 
 
 def pytest_sessionfinish(session, exitstatus):
-    """ whole test run finishes. """
+    """whole test run finishes."""
     send_message(exitstatus)