more linting

smoke-test/.gitignore

@@ -127,6 +127,7 @@ venv.bak/
 .mypy_cache/
 .dmypy.json
 dmypy.json
+.ruff_cache/
 
 # Pyre type checker
 .pyre/

smoke-test/build.gradle

@@ -46,6 +46,7 @@ task yarnInstall(type: YarnTask) {
 }
 
 task installDev(type: Exec) {
+    inputs.file file('pyproject.toml')
     inputs.file file('requirements.txt')
     outputs.file("${venv_name}/.build_install_dev_sentinel")
     commandLine 'bash', '-x', '-c',
@@ -59,13 +60,13 @@ task lint(type: Exec, dependsOn: installDev) {
         "black --check --diff tests/ && " +
         "isort --check --diff tests/ && " +
         "ruff --statistics tests/ && " +
-        "mypy tests/"
+        "mypy tests/ --check-untyped-defs"
 }
 task lintFix(type: Exec, dependsOn: installDev) {
     commandLine 'bash', '-c',
         "source ${venv_name}/bin/activate && set -x && " +
         "black tests/ && " +
         "isort tests/ && " +
         "ruff --fix tests/ && " +
-        "mypy tests/"
+        "mypy tests/ --check-untyped-defs"
 }

smoke-test/pyproject.toml

@@ -16,7 +16,8 @@ requires-python = ">=3.9"
 extend-exclude = '''
 # A regex preceded with ^/ will apply only to files and directories
 # in the root of the project.
-^/tmp
+tmp
+venv
 '''
 include = '\.pyi?$'
 target-version = ['py310']

smoke-test/requirements.txt

@@ -7,7 +7,11 @@ slack-sdk==3.18.1
 aiohttp
 joblib
 pytest-xdist
+# libaries for linting below this
 black==23.7.0
 isort==5.12.0
 mypy==1.5.1
-ruff==0.0.287
+ruff==0.0.287
+# stub version are copied from metadata-ingestion/setup.py and that should be the source of truth
+types-requests>=2.28.11.6,<=2.31.0.3
+types-PyYAML

smoke-test/tests/delete/delete_test.py

@@ -2,8 +2,7 @@
 import os
 
 import pytest
-from datahub.cli.cli_utils import get_aspects_for_entity
-from datahub.cli.ingest_cli import get_session_and_host
+from datahub.cli.cli_utils import get_aspects_for_entity, get_session_and_host
 
 from tests.utils import (
     delete_urns_from_file,

smoke-test/tests/patch/test_datajob_patches.py

@@ -81,10 +81,12 @@ def test_datajob_inputoutput_dataset_patch(wait_for_healthchecks):
 
     with DataHubGraph(DataHubGraphConfig()) as graph:
         graph.emit_mcp(mcpw)
-        inputoutput_lineage_read: DataJobInputOutputClass = graph.get_aspect(
+        inputoutput_lineage_read = graph.get_aspect(
             entity_urn=datajob_urn,
             aspect_type=DataJobInputOutputClass,
         )
+        assert inputoutput_lineage_read is not None
+        assert inputoutput_lineage_read.inputDatasetEdges is not None
         assert (
             inputoutput_lineage_read.inputDatasetEdges[0].destinationUrn
             == other_dataset_urn
@@ -124,6 +126,7 @@ def test_datajob_inputoutput_dataset_patch(wait_for_healthchecks):
             entity_urn=datajob_urn,
             aspect_type=DataJobInputOutputClass,
         )
+        assert inputoutput_lineage_read.inputDatasetEdges is not None
         assert len(inputoutput_lineage_read.inputDatasetEdges) == 1
         assert (
             inputoutput_lineage_read.inputDatasetEdges[0].destinationUrn

smoke-test/tests/patch/test_dataset_patches.py

@@ -171,6 +171,7 @@ def test_field_terms_patch(wait_for_healthchecks):
 
         assert field_info
         assert field_info.description == "This is a test field"
+        assert field_info.glossaryTerms is not None
         assert len(field_info.glossaryTerms.terms) == 1
         assert field_info.glossaryTerms.terms[0].urn == new_term.urn
 
@@ -188,6 +189,7 @@ def test_field_terms_patch(wait_for_healthchecks):
 
         assert field_info
         assert field_info.description == "This is a test field"
+        assert field_info.glossaryTerms is not None
         assert len(field_info.glossaryTerms.terms) == 0
 
 
@@ -231,6 +233,7 @@ def test_field_tags_patch(wait_for_healthchecks):
 
         assert field_info
         assert field_info.description == "This is a test field"
+        assert field_info.globalTags is not None
         assert len(field_info.globalTags.tags) == 1
         assert field_info.globalTags.tags[0].tag == new_tag.tag
 
@@ -249,6 +252,7 @@ def test_field_tags_patch(wait_for_healthchecks):
 
         assert field_info
         assert field_info.description == "This is a test field"
+        assert field_info.globalTags is not None
         assert len(field_info.globalTags.tags) == 1
         assert field_info.globalTags.tags[0].tag == new_tag.tag
 
@@ -266,6 +270,7 @@ def test_field_tags_patch(wait_for_healthchecks):
 
         assert field_info
         assert field_info.description == "This is a test field"
+        assert field_info.globalTags is not None
         assert len(field_info.globalTags.tags) == 0
 
 

smoke-test/tests/privileges/test_privileges.py

@@ -1,6 +1,5 @@
 import pytest
 import tenacity
-from datahub.cli.cli_utils import get_session_login_as
 
 from tests.privileges.utils import (
     assign_role,
@@ -20,6 +19,7 @@
     get_frontend_session,
     get_frontend_url,
     get_sleep_info,
+    login_as,
     wait_for_healthcheck_util,
     wait_for_writes_to_sync,
 )
@@ -158,8 +158,8 @@ def _ensure_can_create_user_policy(session, json):
 @pytest.mark.dependency(depends=["test_healthchecks"])
 def test_privilege_to_create_and_manage_secrets():
     (admin_user, admin_pass) = get_admin_credentials()
-    admin_session = get_session_login_as(admin_user, admin_pass)
-    user_session = get_session_login_as("user", "user")
+    admin_session = login_as(admin_user, admin_pass)
+    user_session = login_as("user", "user")
     secret_urn = "urn:li:dataHubSecret:TestSecretName"
 
     # Verify new user can't create secrets
@@ -213,8 +213,8 @@ def test_privilege_to_create_and_manage_secrets():
 @pytest.mark.dependency(depends=["test_healthchecks"])
 def test_privilege_to_create_and_manage_ingestion_source():
     (admin_user, admin_pass) = get_admin_credentials()
-    admin_session = get_session_login_as(admin_user, admin_pass)
-    user_session = get_session_login_as("user", "user")
+    admin_session = login_as(admin_user, admin_pass)
+    user_session = login_as("user", "user")
 
     # Verify new user can't create ingestion source
     create_ingestion_source = {
@@ -322,8 +322,8 @@ def test_privilege_to_create_and_manage_ingestion_source():
 @pytest.mark.dependency(depends=["test_healthchecks"])
 def test_privilege_to_create_and_manage_access_tokens():
     (admin_user, admin_pass) = get_admin_credentials()
-    admin_session = get_session_login_as(admin_user, admin_pass)
-    user_session = get_session_login_as("user", "user")
+    admin_session = login_as(admin_user, admin_pass)
+    user_session = login_as("user", "user")
 
     # Verify new user can't create access token
     create_access_token = {
@@ -406,8 +406,8 @@ def test_privilege_to_create_and_manage_access_tokens():
 @pytest.mark.dependency(depends=["test_healthchecks"])
 def test_privilege_to_create_and_manage_policies():
     (admin_user, admin_pass) = get_admin_credentials()
-    admin_session = get_session_login_as(admin_user, admin_pass)
-    user_session = get_session_login_as("user", "user")
+    admin_session = login_as(admin_user, admin_pass)
+    user_session = login_as("user", "user")
 
     # Verify new user can't create a policy
     create_policy = {
@@ -502,8 +502,8 @@ def test_privilege_to_create_and_manage_policies():
 @pytest.mark.dependency(depends=["test_healthchecks"])
 def test_privilege_from_group_role_can_create_and_manage_secret():
     (admin_user, admin_pass) = get_admin_credentials()
-    admin_session = get_session_login_as(admin_user, admin_pass)
-    user_session = get_session_login_as("user", "user")
+    admin_session = login_as(admin_user, admin_pass)
+    user_session = login_as("user", "user")
     secret_urn = "urn:li:dataHubSecret:TestSecretName"
 
     # Verify new user can't create secrets

smoke-test/tests/privileges/utils.py

@@ -1,7 +1,5 @@
-from datahub.cli import cli_utils
-
 from tests.consistency_utils import wait_for_writes_to_sync
-from tests.utils import get_admin_credentials, get_frontend_url
+from tests.utils import get_admin_credentials, get_frontend_url, login_as
 
 
 def set_base_platform_privileges_policy_status(status, session):
@@ -164,14 +162,6 @@ def create_user(session, email, password):
     return admin_session
 
 
-def login_as(username, password):
-    return cli_utils.get_session_login_as(
-        username=username,
-        password=password,
-        frontend_url=get_frontend_url(),
-    )
-
-
 def remove_user(session, urn):
     json = {
         "query": """mutation removeUser($urn: String!) {\n

smoke-test/tests/tokens/revokable_access_token_test.py

@@ -1,11 +1,11 @@
 import os
 
 import pytest
-from datahub.cli import cli_utils
 
 from tests.utils import (
     get_admin_credentials,
     get_frontend_url,
+    login_as,
     wait_for_healthcheck_util,
     wait_for_writes_to_sync,
 )
@@ -32,7 +32,7 @@ def test_healthchecks(wait_for_healthchecks):
 @pytest.fixture(scope="class", autouse=True)
 def custom_user_setup():
     """Fixture to execute setup before and tear down after all tests are run"""
-    admin_session = loginAs(admin_user, admin_pass)
+    admin_session = login_as(admin_user, admin_pass)
 
     res_data = removeUser(admin_session, "urn:li:corpuser:user")
     assert res_data
@@ -80,7 +80,7 @@ def custom_user_setup():
 
     # signUp will override the session cookie to the new user to be signed up.
     admin_session.cookies.clear()
-    admin_session = loginAs(admin_user, admin_pass)
+    admin_session = login_as(admin_user, admin_pass)
 
     # Make user created user is there.
     res_data = listUsers(admin_session)
@@ -109,7 +109,7 @@ def custom_user_setup():
 @pytest.fixture(autouse=True)
 def access_token_setup():
     """Fixture to execute asserts before and after a test is run"""
-    admin_session = loginAs(admin_user, admin_pass)
+    admin_session = login_as(admin_user, admin_pass)
 
     res_data = listAccessTokens(admin_session)
     assert res_data
@@ -130,7 +130,7 @@ def access_token_setup():
 
 @pytest.mark.dependency(depends=["test_healthchecks"])
 def test_admin_can_create_list_and_revoke_tokens(wait_for_healthchecks):
-    admin_session = loginAs(admin_user, admin_pass)
+    admin_session = login_as(admin_user, admin_pass)
 
     # Using a super account, there should be no tokens
     res_data = listAccessTokens(admin_session)
@@ -187,7 +187,7 @@ def test_admin_can_create_list_and_revoke_tokens(wait_for_healthchecks):
 
 @pytest.mark.dependency(depends=["test_healthchecks"])
 def test_admin_can_create_and_revoke_tokens_for_other_user(wait_for_healthchecks):
-    admin_session = loginAs(admin_user, admin_pass)
+    admin_session = login_as(admin_user, admin_pass)
 
     # Using a super account, there should be no tokens
     res_data = listAccessTokens(admin_session)
@@ -244,7 +244,7 @@ def test_admin_can_create_and_revoke_tokens_for_other_user(wait_for_healthchecks
 
 @pytest.mark.dependency(depends=["test_healthchecks"])
 def test_non_admin_can_create_list_revoke_tokens(wait_for_healthchecks):
-    user_session = loginAs("user", "user")
+    user_session = login_as("user", "user")
 
     # Normal user should be able to generate token for himself.
     res_data = generateAccessToken_v2(user_session, "urn:li:corpuser:user")
@@ -299,7 +299,7 @@ def test_non_admin_can_create_list_revoke_tokens(wait_for_healthchecks):
 
 @pytest.mark.dependency(depends=["test_healthchecks"])
 def test_admin_can_manage_tokens_generated_by_other_user(wait_for_healthchecks):
-    admin_session = loginAs(admin_user, admin_pass)
+    admin_session = login_as(admin_user, admin_pass)
 
     # Using a super account, there should be no tokens
     res_data = listAccessTokens(admin_session)
@@ -309,7 +309,7 @@ def test_admin_can_manage_tokens_generated_by_other_user(wait_for_healthchecks):
     assert len(res_data["data"]["listAccessTokens"]["tokens"]) == 0
 
     admin_session.cookies.clear()
-    user_session = loginAs("user", "user")
+    user_session = login_as("user", "user")
     res_data = generateAccessToken_v2(user_session, "urn:li:corpuser:user")
     assert res_data
     assert res_data["data"]
@@ -329,7 +329,7 @@ def test_admin_can_manage_tokens_generated_by_other_user(wait_for_healthchecks):
 
     # Admin should be able to list other tokens
     user_session.cookies.clear()
-    admin_session = loginAs(admin_user, admin_pass)
+    admin_session = login_as(admin_user, admin_pass)
     res_data = listAccessTokens(
         admin_session, [{"field": "ownerUrn", "values": ["urn:li:corpuser:user"]}]
     )
@@ -349,7 +349,7 @@ def test_admin_can_manage_tokens_generated_by_other_user(wait_for_healthchecks):
 
     # Admin can delete token created by someone else.
     admin_session.cookies.clear()
-    admin_session = loginAs(admin_user, admin_pass)
+    admin_session = login_as(admin_user, admin_pass)
     res_data = revokeAccessToken(admin_session, user_tokenId)
     assert res_data
     assert res_data["data"]
@@ -360,7 +360,7 @@ def test_admin_can_manage_tokens_generated_by_other_user(wait_for_healthchecks):
 
     # Using a normal account, check that all its tokens where removed.
     user_session.cookies.clear()
-    user_session = loginAs("user", "user")
+    user_session = login_as("user", "user")
     res_data = listAccessTokens(
         user_session, [{"field": "ownerUrn", "values": ["urn:li:corpuser:user"]}]
     )
@@ -370,7 +370,7 @@ def test_admin_can_manage_tokens_generated_by_other_user(wait_for_healthchecks):
     assert len(res_data["data"]["listAccessTokens"]["tokens"]) == 0
 
     # Using the super account, check that all tokens where removed.
-    admin_session = loginAs(admin_user, admin_pass)
+    admin_session = login_as(admin_user, admin_pass)
     res_data = listAccessTokens(
         admin_session, [{"field": "ownerUrn", "values": ["urn:li:corpuser:user"]}]
     )
@@ -382,7 +382,7 @@ def test_admin_can_manage_tokens_generated_by_other_user(wait_for_healthchecks):
 
 @pytest.mark.dependency(depends=["test_healthchecks"])
 def test_non_admin_can_not_generate_tokens_for_others(wait_for_healthchecks):
-    user_session = loginAs("user", "user")
+    user_session = login_as("user", "user")
     # Normal user should not be able to generate token for another user
     res_data = generateAccessToken_v2(user_session, f"urn:li:corpuser:{admin_user}")
     assert res_data
@@ -470,14 +470,6 @@ def revokeAccessToken(session, tokenId):
     return response.json()
 
 
-def loginAs(username, password):
-    return cli_utils.get_session_login_as(
-        username=username,
-        password=password,
-        frontend_url=get_frontend_url(),
-    )
-
-
 def removeUser(session, urn):
     # Remove user
     json = {

smoke-test/tests/utils.py

@@ -19,6 +19,10 @@
 
 def get_frontend_session():
     username, password = get_admin_credentials()
+    return login_as(username, password)
+
+
+def login_as(username: str, password: str):
     return cli_utils.get_session_login_as(
         username=username, password=password, frontend_url=get_frontend_url()
     )