Fix test_python_client.py

clients/python/test_python_client.py

@@ -33,6 +33,10 @@
 import airflow_client.client
 import pytest
 
+from airflow.api_fastapi.app import create_app
+from airflow.auth.managers.simple.datamodels.login import LoginBody
+from airflow.providers.fab.auth_manager.api_fastapi.services.login import FABAuthManagerLogin
+
 try:
     # If you have rich installed, you will have nice colored output of the API responses
     from rich import print
@@ -42,35 +46,27 @@
 from airflow_client.client.api import config_api, dag_api, dag_run_api
 from airflow_client.client.model.dag_run import DAGRun
 
-from airflow.auth.managers.simple.simple_auth_manager import SimpleAuthManager
-from airflow.auth.managers.simple.user import SimpleAuthManagerUser
-
 # The client must use the authentication and authorization parameters
 # in accordance with the API server security policy.
 # Examples for each auth method are provided below, use the example that
 # satisfies your auth use case.
 #
-# In case of the basic authentication below, make sure that Airflow is
-# configured also with the basic_auth as backend additionally to regular session backend needed
-# by the UI. In the `[api]` section of your `airflow.cfg` set:
+# The example below use the default FabAuthManager, in case your airflow api server use a different
+# auth manager for instance AwsAuthManagerUser or SimpleAuthManager make sure to generate the token with
+# appropriate AuthManager.
+# This is defined in the `[api]` section of your `airflow.cfg`:
 #
-# auth_backend = airflow.providers.fab.auth_manager.api.auth.backend.session,airflow.providers.fab.auth_manager.api.auth.backend.basic_auth
+# auth_manager = airflow.auth.managers.simple.simple_auth_manager.SimpleAuthManager
 #
 # Make sure that your user/name are configured properly - using the user/password that has admin
 # privileges in Airflow
 
-# Configure HTTP basic authorization: Basic
-auth_manager = SimpleAuthManager()
+# Used to initialize FAB and the auth manager, necessary for creating the token.
+create_app()
+
+access_token = FABAuthManagerLogin.create_token(LoginBody(username="admin", password="admin")).jwt_token
 configuration = airflow_client.client.Configuration(
     host="http://localhost:8080/public",
-    api_key={
-        "Authorization": "Bearer "
-        + auth_manager._get_token_signer().generate_signed_token(
-            {
-                **auth_manager.serialize_user(SimpleAuthManagerUser(username="test", role="admin")),
-            }
-        )
-    },
 )
 
 # Make sure in the [core] section, the  `load_examples` config is set to True in your airflow.cfg
@@ -81,7 +77,9 @@
 # Enter a context with an instance of the API client
 @pytest.mark.execution_timeout(400)
 def test_python_client():
-    with airflow_client.client.ApiClient(configuration) as api_client:
+    with airflow_client.client.ApiClient(
+        configuration, header_name="Authorization", header_value=f"Bearer {access_token}"
+    ) as api_client:
         errors = False
 
         print("[blue]Getting DAG list")
@@ -90,14 +88,12 @@ def test_python_client():
             try:
                 dag_api_instance = dag_api.DAGApi(api_client)
                 api_response = dag_api_instance.get_dags()
-                print(api_response)
             except airflow_client.client.OpenApiException as e:
                 print(f"[red]Exception when calling DagAPI->get_dags: {e}\n")
                 errors = True
                 time.sleep(6)
                 max_retries -= 1
             else:
-                errors = False
                 print("[green]Getting DAG list successful")
                 break
 
@@ -118,13 +114,13 @@ def test_python_client():
             # dag_run id is generated randomly to allow multiple executions of the script
             dag_run = DAGRun(
                 dag_run_id="some_test_run_" + uuid.uuid4().hex,
+                logical_date=None,
             )
             api_response = dag_run_api_instance.post_dag_run(DAG_ID, dag_run)
             print(api_response)
         except airflow_client.client.exceptions.OpenApiException as e:
             print(f"[red]Exception when calling DAGRunAPI->post_dag_run: {e}\n")
-            # TODO(pierrejeambrun): We need to fix post_dag_run to not return 422
-            errors = False
+            errors = True
         else:
             print("[green]Posting DAG Run successful")
 
@@ -135,7 +131,7 @@ def test_python_client():
             api_response = conf_api_instance.get_config()
             print(api_response)
         except airflow_client.client.OpenApiException as e:
-            if "FORBIDDEN" in str(e):
+            if "Your Airflow administrator chose" in str(e):
                 print(
                     "[yellow]You need to set `expose_config = True` in Airflow configuration"
                     " in order to retrieve configuration."