test: Adjust flow.py file

Techassi · Techassi · commit b40488169815 · 2026-06-16T21:59:25.000+02:00
diff --git a/tests/templates/kuttl/upgrade/03-install-test-nifi.yaml b/tests/templates/kuttl/upgrade/03-install-test-nifi.yaml
@@ -20,3 +20,8 @@ spec:
           image: oci.stackable.tech/sdp/testing-tools/nifi:0.3.0-stackable0.0.0-dev
           stdin: true
           tty: true
+          env:
+            - name: NAMESPACE
+              valueFrom:
+                fieldRef:
+                  fieldPath: metadata.namespace
diff --git a/tests/templates/kuttl/upgrade/04-assert.yaml.j2 b/tests/templates/kuttl/upgrade/04-assert.yaml.j2
@@ -4,6 +4,6 @@ kind: TestAssert
 timeout: 300
 commands:
   - script: kubectl exec -n $NAMESPACE test-nifi-0 -- python /tmp/test_nifi.py -u admin -p supersecretpassword -n $NAMESPACE -c 3
-  - script: kubectl exec -n $NAMESPACE test-nifi-0 -- sh -c "python /tmp/flow.py -e https://test-nifi-node-default-0.test-nifi-node-default-headless.$NAMESPACE.svc.cluster.local:8443 run template /tmp/generate-and-log-flowfiles.xml > /tmp/old_input"
+  - script: kubectl exec -n $NAMESPACE test-nifi-0 -- sh -c "python /tmp/flow.py > /tmp/old_input"
   # This tests if the output contains an Error or zero flow files are queued, which also indicates that something went wrong
   - script: kubectl exec -n $NAMESPACE test-nifi-0 -- sh -c "cat /tmp/old_input | grep -Eov 'Error|\b0\b'"
diff --git a/tests/templates/kuttl/upgrade/04-prepare-test-nifi.yaml.j2 b/tests/templates/kuttl/upgrade/04-prepare-test-nifi.yaml.j2
@@ -5,5 +5,5 @@ commands:
   - script: kubectl cp -n $NAMESPACE ./test_nifi_metrics.py  test-nifi-0:/tmp
   - script: kubectl cp -n $NAMESPACE ./test_nifi.py  test-nifi-0:/tmp
   - script: kubectl cp -n $NAMESPACE ./flow.py  test-nifi-0:/tmp
-  - script: kubectl cp -n $NAMESPACE ./generate-and-log-flowfiles.xml  test-nifi-0:/tmp
+  - script: kubectl cp -n $NAMESPACE ./generate-and-log-flowfiles.json test-nifi-0:/tmp
   - script: kubectl cp -n $NAMESPACE ./cacert.pem  test-nifi-0:/tmp
diff --git a/tests/templates/kuttl/upgrade/flow.py b/tests/templates/kuttl/upgrade/flow.py
@@ -1,6 +1,4 @@
 #!/usr/bin/env python3
-# This script is used to create a flow from a template, run it for a couple
-# of seconds, stop it and then query the number of queued flow files.
 
 import nipyapi
 from nipyapi.canvas import (
@@ -11,25 +9,19 @@
     recurse_flow,
 )
 from nipyapi.security import service_login
-from nipyapi.templates import get_template, upload_template, deploy_template
-from nipyapi.nifi.models import FlowEntity
+from nipyapi.nifi.api_client import ApiClient
 
-from typing import Union, List
+from typing import Union
 from dataclasses import dataclass
 
-import re
 import os
-import sys
 import time
-import argparse
 import urllib3
 
+ENDPOINT = f"https://test-nifi-node-default-0.test-nifi-node-default-headless.{os.environ['NAMESPACE']}.svc.cluster.local:8443"
 USERNAME = "admin"
-PASSWORD = "supersecretpassword"
-TEMPLATE_NAME = "generate-and-log-flowfiles"
-TEMPLATE_FILE = f"{TEMPLATE_NAME}.xml"
-JSON_NAME = "generate-and-log-flowfiles"
-JSON_FILE = f"{JSON_NAME}.json"
+PASSWORD = "supersecretpassword"  # Can we use open("/nifi-users/admin").read() here?
+FLOWFILE = "/tmp/generate-and-log-flowfiles.json"
 CONNECTION_NAME = "output-connection"
 
 
@@ -38,170 +30,83 @@ class Error:
     message: str
 
 
-def parse_args(args: List[str]) -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Setup a flow and query the status.")
-    parser.add_argument(
-        "-e",
-        "--endpoint",
-        help="Nifi endpoint URL (must be https). As of 2022-08-29 we cant use https://nifi:8443 here because Nifi's check for the host header validity fails. The error message looks like: <h2>The request contained an invalid host header [<code>nifi:8443</code>] in the request [<code>/nifi-api</code>]. Check for request manipulation or third-party intercept.</h2>",
-        required=True,
-    )
-    parser.add_argument(
-        "-u",
-        "--user",
-        help="Nifi login user.",
-    )
-    parser.add_argument(
-        "-p",
-        "--password",
-        help="Nifi login password.",
-    )
-    subparsers = parser.add_subparsers(dest="subcommand", required=True)
-    run_parser = subparsers.add_parser(
-        "run",
-        help="Install and run a flow.",
-    )
-    format_parser = run_parser.add_subparsers(dest="formatcommand")
-    template_parser = format_parser.add_parser(
-        "template", help="Use the template format for the flow."
-    )
-    template_parser.add_argument(
-        "template",
-        help="Nifi template file.",
-    )
-    json_parser = format_parser.add_parser(
-        "json", help="Use the json format for the flow."
-    )
-    json_parser.add_argument(
-        "json",
-        help="Nifi json file.",
-    )
-    subparsers.add_parser(
-        "query",
-        help="Query the flow.",
-    )
-    return parser.parse_args(args)
-
-
-def login(endpoint: str, user: str, passwd: str) -> bool:
-    nipyapi.config.nifi_config.host = f"{endpoint}/nifi-api"
-    nipyapi.config.nifi_config.verify_ssl = False
-    return service_login(username=user, password=passwd)
-
-
-def flow_from_template(
-    pg_id: str, template_file: str, template_name: str
-) -> FlowEntity:
-    upload_template(pg_id, template_file)
-    template_id = get_template(template_name).id
-    return deploy_template(pg_id, template_id, 200, 0)
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
 
 def schedule(pg_id: str, toggle: bool) -> bool:
     for controller in list_all_controllers(pg_id):
         schedule_controller(controller, scheduled=toggle)
+
     return schedule_process_group(pg_id, scheduled=toggle)
 
 
 def flow_files_queued(connection_name: str) -> Union[str, Error]:
     """Returns the current input record number for the given component"""
     flow = recurse_flow()
-    c = [
-        c
-        for c in flow.process_group_flow.flow.connections
-        if c.status.name == connection_name
-    ]
-    if c:
-        return c[0].status.aggregate_snapshot.flow_files_queued
-    else:
-        # for nifi 2.x the flow is loaded via api into a new process group
-        # we need to look for the connection in that process group instead
-        if flow.process_group_flow.flow.process_groups:
-            inner_flow = recurse_flow(flow.process_group_flow.flow.process_groups[0].id)
-            inner_c = [
-                inner_c
-                for inner_c in inner_flow.process_group_flow.flow.connections
-                if inner_c.status.name == connection_name
-            ]
-            if inner_c:
-                return inner_c[0].status.aggregate_snapshot.flow_files_queued
-            else:
-                return Error(f"No connection named '{connection_name}' found.")
+
+    if flow.process_group_flow.flow.process_groups:
+        inner_flow = recurse_flow(flow.process_group_flow.flow.process_groups[0].id)
+        conn = [
+            conn
+            for conn in inner_flow.process_group_flow.flow.connections
+            if conn.status.name == connection_name
+        ]
+        if conn:
+            return conn[0].status.aggregate_snapshot.flow_files_queued
         else:
             return Error(f"No connection named '{connection_name}' found.")
+    else:
+        return Error(f"No connection named '{connection_name}' found.")
 
 
 def main():
-    # turn this on to debug Nifi api calls. Also need to import logging
-    # logging.getLogger().setLevel(logging.INFO)
-
-    args = parse_args(sys.argv[1:])
-
-    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
-
-    username = args.user or USERNAME
-    password = args.password or PASSWORD
-
-    if args.subcommand == "run":
-        login(args.endpoint, username, password)
-        pg_id = get_root_pg_id()
-
-        if args.formatcommand == "template":
-            template_file = args.template or TEMPLATE_FILE
-            template_name = re.sub(r"\..+$", "", os.path.basename(template_file))
-
-            flow_from_template(pg_id, template_file, template_name)
-
-        elif args.formatcommand == "json":
-            json_file = args.json or JSON_FILE
-            json_name = re.sub(r"\..+$", "", os.path.basename(json_file))
-
-            if not nipyapi.config.nifi_config.api_client:
-                nipyapi.config.nifi_config.api_client = nipyapi.nifi.ApiClient()
-
-            header_params = {}
-            header_params["Accept"] = (
-                nipyapi.config.nifi_config.api_client.select_header_accept(
-                    ["application/json"]
-                )
-            )
-            header_params["Content-Type"] = (
-                nipyapi.config.nifi_config.api_client.select_header_content_type(
-                    ["multipart/form-data"]
-                )
-            )
-
-            nipyapi.config.nifi_config.api_client.call_api(
-                "/process-groups/{pg_id}/process-groups/upload",
-                "POST",
-                path_params={"pg_id": pg_id},
-                header_params=header_params,
-                _return_http_data_only=True,
-                post_params=[
-                    ("id", pg_id),
-                    ("groupName", json_name),
-                    ("positionX", 100),
-                    ("positionY", 10),
-                    ("clientId", nipyapi.nifi.FlowApi().generate_client_id()),
-                ],
-                files={"file": json_file},
-                auth_settings=["tokenAuth"],
-            )
-
-            # get the process group id of the newly created process group to pass it to the schedule() function
-            # for nifi 2.x we need to schedule the inner process group and instead of the root process group
-            flow = recurse_flow()
-            inner_flow_id = flow.process_group_flow.flow.process_groups[0].id
-            pg_id = inner_flow_id
-
-        schedule(pg_id, True)  # start
-        time.sleep(5)  # give the flow 5 seconds to run
-        schedule(pg_id, False)  # stop
-        print(flow_files_queued(CONNECTION_NAME))
-
-    elif args.subcommand == "query":
-        login(args.endpoint, username, password)
-        print(flow_files_queued(CONNECTION_NAME))
+    nipyapi.config.nifi_config.host = f"{ENDPOINT}/nifi-api"
+    nipyapi.config.nifi_config.verify_ssl = False
+
+    print(f"Logging in as {USERNAME}")
+    service_login(username=USERNAME, password=PASSWORD)
+    print("Logged in")
+
+    pg_id = get_root_pg_id()
+    print(f"Got root process group id: {pg_id}")
+
+    if not nipyapi.config.nifi_config.api_client:
+        nipyapi.config.nifi_config.api_client = ApiClient()
+
+    header_params = {}
+    header_params["Accept"] = (
+        nipyapi.config.nifi_config.api_client.select_header_accept(["application/json"])
+    )
+    header_params["Content-Type"] = (
+        nipyapi.config.nifi_config.api_client.select_header_content_type(
+            ["multipart/form-data"]
+        )
+    )
+
+    print("Uploading process group")
+    nipyapi.config.nifi_config.api_client.call_api(
+        "/process-groups/{pg_id}/process-groups/upload",
+        "POST",
+        path_params={"pg_id": pg_id},
+        header_params=header_params,
+        _return_http_data_only=True,
+        post_params=[
+            ("id", pg_id),
+            ("groupName", "Upgrade Test"),
+            ("positionX", 100),
+            ("positionY", 10),
+            ("clientId", nipyapi.nifi.FlowApi().generate_client_id()),
+        ],
+        files={"file": FLOWFILE},
+        auth_settings=["tokenAuth"],
+    )
+    print("Process group uploaded")
+
+    schedule(pg_id, True)  # Start the flow
+    time.sleep(5)  # Let the flow run for 5 seconds
+    schedule(pg_id, False)  # Stop the flow
+
+    print(flow_files_queued(CONNECTION_NAME))
 
 
 if __name__ == "__main__":
