diff --git a/README.md b/README.md index fee2dff..0c5acd2 100644 --- a/README.md +++ b/README.md @@ -31,3 +31,24 @@ Or if it is a commandline tool then you might put some example commands here: ``` python -m python_interface_to_workflows --version ``` + +# Using Copier +```bash +mkdir new_directory_path +cd new_directory_path +git init +git remote add origin {origin ssh} +git branch -M main + +cd .. +``` +then either: +```bash +git clone git@github.com:DiamondLightSource/python-interface-to-workflows.git +copier copy {this_repo's_path} {new_directory_path} +``` +or: +```bash +copier copy git@github.com:DiamondLightSource/python-copier-template.git new directory path +``` +rebuild in dev container without cache diff --git a/scripts/checkyamlcompliance.py b/scripts/checkyamlcompliance.py index c862655..834655d 100644 --- a/scripts/checkyamlcompliance.py +++ b/scripts/checkyamlcompliance.py @@ -26,6 +26,10 @@ match yamldata["kind"]: case "Workflow": clst_tmpt = True + case "WorkflowTemplate": + clst_tmpt = True + case "ClusterWorkflowTemplate": + clst_tmpt = True case _: clst_tmpt = False else: diff --git a/scripts/lintyaml.sh b/scripts/lintyaml.sh index 3de474e..566462a 100644 --- a/scripts/lintyaml.sh +++ b/scripts/lintyaml.sh @@ -1,9 +1,10 @@ #!/bin/bash cd src/python_interface_to_workflows/templates -for file in * -do -argo lint "$file" --offline -SUCCESSFULLINT=$? -[ $SUCCESSFULLINT -ne 0 ] && exit 1 +for file in *.txt; do + yaml_file="${file}.yaml" + + mv "$file" "$yaml_file" + argo lint "$yaml_file" --offline + mv "$yaml_file" "$file" done exit 0 diff --git a/scripts/makecopiercorrect.sh b/scripts/makecopiercorrect.sh index 0f2a966..973bd00 100644 --- a/scripts/makecopiercorrect.sh +++ b/scripts/makecopiercorrect.sh @@ -7,16 +7,36 @@ cp ../templates/*example*.txt "../../copier_template/src/{{ project_name }}/temp for file in "../../copier_template/src/{{ project_name }}/workflow_definitions"/* do [[ $file == *.jinja ]] && continue + sed -i 's/python-interface-to-workflows/{{repo_name}}/g' "$file" sed -i 's/DiamondLightSource/{{github_org}}/g' "$file" + + sed -i '1i{% raw %}' "$file" + echo '{% endraw %}' >> "$file" + + sed -i \ + -e 's/{{repo_name}}/{% endraw %}{{repo_name}}{% raw %}/g' \ + -e 's/{{github_org}}/{% endraw %}{{github_org}}{% raw %}/g' \ + "$file" + mv "$file" "$file.jinja" done for file in "../../copier_template/src/{{ project_name }}/templates"/* do [[ $file == *.jinja ]] && continue + sed -i 's/python-interface-to-workflows/{{repo_name}}/g' "$file" sed -i 's/DiamondLightSource/{{github_org}}/g' "$file" + + sed -i '1i{% raw %}' "$file" + echo '{% endraw %}' >> "$file" + + sed -i \ + -e 's/{{repo_name}}/{% endraw %}{{repo_name}}{% raw %}/g' \ + -e 's/{{github_org}}/{% endraw %}{{github_org}}{% raw %}/g' \ + "$file" + mv "$file" "$file.jinja" done git add "../../copier_template/src/{{ project_name }}/workflow_definitions"/* diff --git a/scripts/runthefiles.sh b/scripts/runthefiles.sh index a7453b4..960a7d8 100644 --- a/scripts/runthefiles.sh +++ b/scripts/runthefiles.sh @@ -5,4 +5,5 @@ do uv run "$file" done mv *.yaml ../templates/ +mv *.txt ../templates/ git add -u ../templates/ diff --git a/src/copier_template/pyproject.toml.jinja b/src/copier_template/pyproject.toml.jinja index cad7d38..9eec731 100644 --- a/src/copier_template/pyproject.toml.jinja +++ b/src/copier_template/pyproject.toml.jinja @@ -68,7 +68,6 @@ reportMissingImports = false # Ignore missing stubs in imported modules reportMissingTypeStubs = "none" reportMissingModuleSource="warning" reportInvalidTypeForm="warning" -reportMissingImports="warning" reportUndefinedVariable="warning" reportAbstractUsage="warning" reportArgumentType="warning" diff --git a/src/copier_template/scripts/lintyaml.sh.jinja b/src/copier_template/scripts/lintyaml.sh.jinja index 7b0b6e8..5133bb8 100644 --- a/src/copier_template/scripts/lintyaml.sh.jinja +++ b/src/copier_template/scripts/lintyaml.sh.jinja @@ -1,9 +1,10 @@ #!/bin/bash cd src/{{project_name}}/templates -for file in * -do -argo lint "$file" --offline -SUCCESSFULLINT=$? -[ $SUCCESSFULLINT -ne 0 ] && exit 1 +for file in *.txt; do + yaml_file="${file}.yaml" + + mv "$file" "$yaml_file" + argo lint "$yaml_file" --offline + mv "$yaml_file" "$file" done exit 0 diff --git a/src/copier_template/scripts/runthefiles.sh.jinja b/src/copier_template/scripts/runthefiles.sh.jinja index ca0bea2..f324322 100644 --- a/src/copier_template/scripts/runthefiles.sh.jinja +++ b/src/copier_template/scripts/runthefiles.sh.jinja @@ -5,4 +5,5 @@ do uv run "$file" done mv *.yaml ../templates/ +mv *.txt ../templates/ git add -u ../templates/ diff --git a/src/copier_template/src/README.md.jinja b/src/copier_template/src/README.md.jinja index 00f1d90..e893d23 100644 --- a/src/copier_template/src/README.md.jinja +++ b/src/copier_template/src/README.md.jinja @@ -4,7 +4,10 @@ HOST=https://argo-workflows.workflows.diamond.ac.uk/ (to submit to the production cluster) DEFAULT_IMAGE= (usually python 3.10) -NAMESPACE= (the cluster you wish to run the templates on) +VISIT= (the Visit you wish to run the template on) +TOKEN= +EXPIRY= +AUTH= 3. Build the dev container @@ -22,3 +25,20 @@ submit_workflow_to_argo(w) NOTE: Be sure to remove this line upon commiting changes, as all workflow definition files are by default, ran on pre-commit, to ensure that any yaml files they create are up to date. + +# Building a custom image +While in src, the same folder as a Dockerfile: + +podman build -t ghcr.io/Your-Github-Name/image-name . +podman login ghcr.io +podman push ghcr.io/Your-Github-Name/image-name + +Then go to your github profile, packages, and set image-name's visibility to public +After this, you may add 'image' in the script decorator, to run specific scripts within that image +Alternatively, you can set the default image at the top of the file by adding: + +```python +global_config.set_class_defaults( # pyright: ignore + Script, image=str(os.environ.get("DEFAULT_IMAGE")) +) +``` diff --git a/src/copier_template/src/{{ project_name }}/auth/keycloak_checker.py.jinja b/src/copier_template/src/{{ project_name }}/auth/keycloak_checker.py.jinja index 04dd688..ff9db43 100644 --- a/src/copier_template/src/{{ project_name }}/auth/keycloak_checker.py.jinja +++ b/src/copier_template/src/{{ project_name }}/auth/keycloak_checker.py.jinja @@ -1,20 +1,30 @@ import os +from typing import TypedDict, cast import dotenv from keycloak import KeycloakOpenID from keycloak.pkce_utils import generate_code_challenge, generate_code_verifier -from {{project_name}}.auth.open_auth_url import open_auth_url +from {{project_name}}.auth.open_auth_url import token_expired + + +class TokenResponse(TypedDict): + access_token: str + refresh_token: str + + +class DecodedToken(TypedDict): + exp: int def set_token_env_variable() -> str: keycloak_openid = KeycloakOpenID( - client_id="workflows-dashboard", - server_url="https://identity.diamond.ac.uk/", - realm_name="dls", - client_secret_key="", - pool_maxsize=1, - ) + client_id="workflows-cli", + server_url="https://identity.diamond.ac.uk/", + realm_name="dls", + client_secret_key="", + pool_maxsize=1, + ) port = 8000 code_verifier = generate_code_verifier() code_challenge, code_challenge_method = generate_code_challenge(code_verifier) @@ -25,14 +35,38 @@ def set_token_env_variable() -> str: code_challenge=code_challenge, code_challenge_method=code_challenge_method, ) - open_auth_url(auth_url) - token: dict[str, str] = ( # pyright: ignore[reportUnknownVariableType] - keycloak_openid.token( # pyright: ignore[reportUnknownMemberType] - grant_type="authorization_code", - code=os.environ["AUTH"], - redirect_uri=f"http://localhost:{port}/", - code_verifier=code_verifier, + if token_expired(auth_url, port): + token = cast( + TokenResponse, + keycloak_openid.token( # pyright: ignore[reportUnknownMemberType] + grant_type="authorization_code", + code=str(os.environ.get("AUTH")), + redirect_uri=f"http://localhost:{port}/", + code_verifier=code_verifier, + ), + ) + else: + token = cast( + TokenResponse, + keycloak_openid.refresh_token( # pyright: ignore[reportUnknownMemberType] + str(os.environ.get("REFRESHTOKEN")) + ), ) + token_info = cast( + DecodedToken, + keycloak_openid.decode_token( # pyright: ignore[reportUnknownMemberType] + token["access_token"] + ), ) - dotenv.set_key("src/.env", "TOKEN", token["access_token"].strip("'")) - return token["access_token"] # pyright: ignore[reportUnknownArgumentType] + try: + expire_time = int(token_info["exp"]) + 1800 + dotenv.set_key("src/.env", "EXPIRY", str(expire_time)) + dotenv.set_key("src/.env", "TOKEN", token["access_token"].strip("'")) + dotenv.set_key("src/.env", "REFRESHTOKEN", token["refresh_token"].strip("'")) + except AttributeError: + print("ERROR:") + exit(1) + finally: + dotenv.load_dotenv(dotenv_path="src/.env", override=True) + + return token["access_token"] diff --git a/src/copier_template/src/{{ project_name }}/auth/open_auth_url.py b/src/copier_template/src/{{ project_name }}/auth/open_auth_url.py index 5a18266..f8c2581 100644 --- a/src/copier_template/src/{{ project_name }}/auth/open_auth_url.py +++ b/src/copier_template/src/{{ project_name }}/auth/open_auth_url.py @@ -1,10 +1,13 @@ import os import socket +import time import urllib.parse import webbrowser from http.server import BaseHTTPRequestHandler, HTTPServer from typing import cast +import dotenv + class _ReusingHTTPServer(HTTPServer): allow_reuse_address = True @@ -26,12 +29,23 @@ def do_GET(self): self.wfile.write(b"Missing authorization code.") -def open_auth_url(auth_url: str) -> None: - httpd = _ReusingHTTPServer(("localhost", 8000), CallbackHandler) +def token_expired(auth_url: str, port: int) -> bool: + dotenv.load_dotenv(dotenv_path="src/.env", override=True) + expiry_str: str = str(os.environ.get("EXPIRY")).strip("'") + if (expiry_str == "" or int(expiry_str)) <= float(time.time()): + _open_auth_url(auth_url, port) + return True + else: + return False + + +def _open_auth_url(auth_url: str, port: int) -> None: + httpd = _ReusingHTTPServer(("localhost", port), CallbackHandler) webbrowser.open(auth_url) try: httpd.handle_request() os.environ["AUTH"] = httpd.auth_code + dotenv.set_key("src/.env", "AUTH", httpd.auth_code) except OSError: os.environ["AUTH"] = "" print("ERROR: Port in use. Please restart your terminal.") diff --git a/src/copier_template/src/{{ project_name }}/submit_to_argo.py.jinja b/src/copier_template/src/{{ project_name }}/submit_to_argo.py.jinja index f0e49e4..9a0ea54 100644 --- a/src/copier_template/src/{{ project_name }}/submit_to_argo.py.jinja +++ b/src/copier_template/src/{{ project_name }}/submit_to_argo.py.jinja @@ -12,12 +12,11 @@ from {{project_name}}.auth.keycloak_checker import set_token_env_variable def submit_workflow_to_argo(w: Workflow): set_token_env_variable() dotenv.load_dotenv(dotenv_path="src/.env", override=True) - w.namespace = os.environ.get("NAMESPACE") + w.namespace = os.environ.get("VISIT") w.workflows_service = WorkflowsService( host=str(os.environ.get("HOST")).strip("'"), token=str(os.environ.get("TOKEN")).strip("'"), ) - global_config.image = str(os.environ.get("DEFAULT_IMAGE")) submitted_w = cast(m.Workflow, w.create()) name = submitted_w.metadata.name namespace = submitted_w.metadata.namespace diff --git a/src/copier_template/src/{{ project_name }}/templates/example.txt.jinja b/src/copier_template/src/{{ project_name }}/templates/example.txt.jinja index 736759b..387db1c 100644 --- a/src/copier_template/src/{{ project_name }}/templates/example.txt.jinja +++ b/src/copier_template/src/{{ project_name }}/templates/example.txt.jinja @@ -1,5 +1,6 @@ +{% raw %} apiVersion: argoproj.io/v1alpha1 -kind: Workflow +kind: WorkflowTemplate metadata: generateName: hera-example- annotations: @@ -7,7 +8,7 @@ metadata: Replicates the functionality of example.yaml workflows.argoproj.io/title: example remade via hera - workflows.diamond.ac.uk/repository: https://github.com/{{github_org}}/{{repo_name}} + workflows.diamond.ac.uk/repository: https://github.com/{% endraw %}{{github_org}}{% raw %}/{% endraw %}{{repo_name}}{% raw %} labels: workflows.diamond.ac.uk/science-group-examples: 'true' spec: @@ -212,3 +213,4 @@ spec: resources: requests: storage: 1Gi +{% endraw %} diff --git a/src/copier_template/src/{{ project_name }}/templates/example_in_image.txt.jinja b/src/copier_template/src/{{ project_name }}/templates/example_in_image.txt.jinja new file mode 100644 index 0000000..78b1316 --- /dev/null +++ b/src/copier_template/src/{{ project_name }}/templates/example_in_image.txt.jinja @@ -0,0 +1,198 @@ +{% raw %} +apiVersion: argoproj.io/v1alpha1 +kind: Workflow +metadata: + generateName: hera-example-in-image- + annotations: + workflows.argoproj.io/description: |- + Replicates the functionality of + example.yaml + workflows.argoproj.io/title: example remade via hera + workflows.diamond.ac.uk/repository: https://github.com/{% endraw %}{{github_org}}{% raw %}/{% endraw %}{{repo_name}}{% raw %} + labels: + workflows.diamond.ac.uk/science-group-examples: 'true' +spec: + entrypoint: workflowentry + templates: + - name: workflowentry + dag: + tasks: + - name: params + template: generate-parameters + arguments: + parameters: + - name: png + value: 'True' + - name: jpg + value: 'True' + - name: jpeg + value: 'True' + - name: tif + value: 'True' + - name: tiff + value: 'True' + - name: create-image + depends: params + template: create-image + withParam: '{{tasks.params.outputs.parameters.out-parameters}}' + arguments: + parameters: + - name: width + value: '{{item.width}}' + - name: height + value: '{{item.height}}' + - name: weights + value: '{{item.weights}}' + - name: extension + value: '{{item.extension}}' + - name: to-hdf5 + depends: create-image + template: to-hdf5 + arguments: + parameters: + - name: paths + value: '{{tasks.create-image.outputs.parameters.out-paths}}' + - name: generate-parameters + inputs: + parameters: + - name: png + - name: jpg + - name: jpeg + - name: tif + - name: tiff + outputs: + parameters: + - name: out-parameters + valueFrom: + path: /tmp/parameters.json + script: + image: ghcr.io/matt-carre/{% endraw %}{{repo_name}}{% raw %}-default-image + source: |- + import os + import sys + sys.path.append(os.getcwd()) + import json + try: jpeg = json.loads(r'''{{inputs.parameters.jpeg}}''') + except: jpeg = r'''{{inputs.parameters.jpeg}}''' + try: jpg = json.loads(r'''{{inputs.parameters.jpg}}''') + except: jpg = r'''{{inputs.parameters.jpg}}''' + try: png = json.loads(r'''{{inputs.parameters.png}}''') + except: png = r'''{{inputs.parameters.png}}''' + try: tif = json.loads(r'''{{inputs.parameters.tif}}''') + except: tif = r'''{{inputs.parameters.tif}}''' + try: tiff = json.loads(r'''{{inputs.parameters.tiff}}''') + except: tiff = r'''{{inputs.parameters.tiff}}''' + + import json + params: list[dict[str, int | list[int] | str] | None] = [{'width': 500, 'height': 500, 'weights': [255, 1, 100], 'extension': 'png'} if png.lower() == 'true' else None, {'width': 600, 'height': 200, 'weights': [100, 150, 100], 'extension': 'jpg'} if jpg.lower() == 'true' else None, {'width': 300, 'height': 400, 'weights': [100, 150, 100], 'extension': 'jpeg'} if jpeg.lower() == 'true' else None, {'width': 300, 'height': 200, 'weights': [230, 100, 1], 'extension': 'tif'} if tif.lower() == 'true' else None, {'width': 200, 'height': 300, 'weights': [230, 100, 1], 'extension': 'tiff'} if tiff.lower() == 'true' else None] + params_to_write: list[dict[str, int | list[int] | str]] = [image_params for image_params in params if image_params is not None] + with open('/tmp/parameters.json', 'w') as f: + json.dump(params_to_write, f) + command: + - python + volumeMounts: + - name: tmpdir + mountPath: /tmp + - name: create-image + inputs: + parameters: + - name: width + - name: height + - name: weights + - name: extension + outputs: + artifacts: + - name: '{{inputs.parameters.extension}}-image' + path: /tmp/{{inputs.parameters.extension}}-image.{{inputs.parameters.extension}} + archive: + none: {} + parameters: + - name: out-paths + valueFrom: + path: /tmp/{{inputs.parameters.extension}}-path.json + script: + image: ghcr.io/matt-carre/{% endraw %}{{repo_name}}{% raw %}-default-image + source: |- + import os + import sys + sys.path.append(os.getcwd()) + import json + try: extension = json.loads(r'''{{inputs.parameters.extension}}''') + except: extension = r'''{{inputs.parameters.extension}}''' + try: height = json.loads(r'''{{inputs.parameters.height}}''') + except: height = r'''{{inputs.parameters.height}}''' + try: weights = json.loads(r'''{{inputs.parameters.weights}}''') + except: weights = r'''{{inputs.parameters.weights}}''' + try: width = json.loads(r'''{{inputs.parameters.width}}''') + except: width = r'''{{inputs.parameters.width}}''' + + import json + from PIL import Image + + def create_pattern(width: int, height: int, weights: tuple[int, int, int]) -> Image.Image: + print(f'width: {width}') + print(f'height: {height}') + print(f'RBG weights: {weights}') + image = Image.new('RGB', (width, height)) + pixels = image.load() + for i in range(width): + for j in range(height): + pixels[i, j] = ((i + j * 50) % weights[0], weights[1], (i * 300 + j) % weights[2]) + return image + image = create_pattern(width, height, weights) + path = f'/tmp/{extension}-image.{extension}' + image.save(path) + with open(f'/tmp/{extension}-path.json', 'w') as f: + json.dump(path, f) + command: + - python + volumeMounts: + - name: tmpdir + mountPath: /tmp + - name: to-hdf5 + inputs: + parameters: + - name: paths + outputs: + artifacts: + - name: hdf5output + path: /tmp/images.hdf5 + archive: + none: {} + script: + image: ghcr.io/matt-carre/{% endraw %}{{repo_name}}{% raw %}-default-image + source: |- + import os + import sys + sys.path.append(os.getcwd()) + import json + try: paths = json.loads(r'''{{inputs.parameters.paths}}''') + except: paths = r'''{{inputs.parameters.paths}}''' + + import h5py + import numpy as np + from PIL import Image + print('creating hdf5 file') + with h5py.File('/tmp/images.hdf5', 'w') as f: + for i, path in enumerate(paths): + path = path.strip('"') + print(f'Got {path}') + with Image.open(path) as image: + arr = np.array(image) + f.create_dataset(f'image_{i}', data=arr, dtype=arr.dtype) + print('done') + command: + - python + volumeMounts: + - name: tmpdir + mountPath: /tmp + volumeClaimTemplates: + - metadata: + name: tmpdir + spec: + accessModes: + - ReadWriteOnce + resources: + requests: + storage: 1Gi +{% endraw %} diff --git a/src/copier_template/src/{{ project_name }}/workflow_definitions/create_example_template.py.jinja b/src/copier_template/src/{{ project_name }}/workflow_definitions/create_example_template.py.jinja index 6a6aa90..42175cf 100644 --- a/src/copier_template/src/{{ project_name }}/workflow_definitions/create_example_template.py.jinja +++ b/src/copier_template/src/{{ project_name }}/workflow_definitions/create_example_template.py.jinja @@ -1,3 +1,4 @@ +{% raw %} from hera.workflows import ( DAG, Artifact, @@ -147,13 +148,13 @@ with Workflow( generate_name="hera-example-", # when running on graphql this should be name entrypoint="workflowentry", api_version="argoproj.io/v1alpha1", - kind="Workflow", # ClusterWorkflowTemplate", when on graphql + kind="WorkflowTemplate", # ClusterWorkflowTemplate", when on graphql labels={"workflows.diamond.ac.uk/science-group-examples": "true"}, annotations={ "workflows.argoproj.io/title": "example remade via hera", "workflows.argoproj.io/description": """Replicates the functionality of example.yaml""", - "workflows.diamond.ac.uk/repository": "https://github.com/{{github_org}}/{{repo_name}}", + "workflows.diamond.ac.uk/repository": "https://github.com/{% endraw %}{{github_org}}{% raw %}/{% endraw %}{{repo_name}}{% raw %}", }, volumes=Volume(name="tmpdir", mount_path="/tmp/", size="1Gi"), ) as w: @@ -178,5 +179,6 @@ example.yaml""", [install, params] >> makeimages >> makehdf5 # pyright: ignore -with open("example.yaml", "w") as div: +with open("example.txt", "w") as div: div.write(w.to_yaml()) # pyright: ignore[reportUnknownMemberType] +{% endraw %} diff --git a/src/copier_template/src/{{ project_name }}/workflow_definitions/create_example_template_within_default_image.py.jinja b/src/copier_template/src/{{ project_name }}/workflow_definitions/create_example_template_within_default_image.py.jinja new file mode 100644 index 0000000..80e5e29 --- /dev/null +++ b/src/copier_template/src/{{ project_name }}/workflow_definitions/create_example_template_within_default_image.py.jinja @@ -0,0 +1,174 @@ +{% raw %} +import os + +from hera.shared import global_config +from hera.workflows import ( + DAG, + Artifact, + Parameter, + Script, + Volume, + Workflow, + script, # pyright: ignore[reportUnknownVariableType] +) +from hera.workflows import models as m +from hera.workflows.archive import NoneArchiveStrategy + +global_config.set_class_defaults( # pyright: ignore + Script, image=str(os.environ.get("DEFAULT_IMAGE")) +) + + +@script( + command=["python"], + outputs=Parameter( + name="out-parameters", value_from=m.ValueFrom(path="/tmp/parameters.json") + ), + volume_mounts=[m.VolumeMount(name="tmpdir", mount_path="/tmp")], +) +def generate_parameters( + png: str, + jpg: str, + jpeg: str, + tif: str, + tiff: str, +): + import json + + params: list[dict[str, int | list[int] | str] | None] = [ + {"width": 500, "height": 500, "weights": [255, 1, 100], "extension": "png"} + if png.lower() == "true" + else None, + {"width": 600, "height": 200, "weights": [100, 150, 100], "extension": "jpg"} + if jpg.lower() == "true" + else None, + {"width": 300, "height": 400, "weights": [100, 150, 100], "extension": "jpeg"} + if jpeg.lower() == "true" + else None, + {"width": 300, "height": 200, "weights": [230, 100, 1], "extension": "tif"} + if tif.lower() == "true" + else None, + {"width": 200, "height": 300, "weights": [230, 100, 1], "extension": "tiff"} + if tiff.lower() == "true" + else None, + ] + params_to_write: list[dict[str, int | list[int] | str]] = [ + image_params for image_params in params if image_params is not None + ] + with open("/tmp/parameters.json", "w") as f: + json.dump(params_to_write, f) + + +@script( + command=["python"], + volume_mounts=[m.VolumeMount(name="tmpdir", mount_path="/tmp")], + outputs=[ + Parameter( + name="out-paths", + value_from=m.ValueFrom( + path="/tmp/{{inputs.parameters.extension}}-path.json" + ), + ), + Artifact( + name="{{inputs.parameters.extension}}-image", + path="/tmp/{{inputs.parameters.extension}}-image.{{inputs.parameters.extension}}", + archive=NoneArchiveStrategy(), + ), + ], +) +def create_image( + width: int, height: int, weights: tuple[int, int, int], extension: str +): + import json + + from PIL import Image + + def create_pattern( + width: int, + height: int, + weights: tuple[int, int, int], + ) -> Image.Image: + print(f"width: {width}") + print(f"height: {height}") + print(f"RBG weights: {weights}") + image = Image.new("RGB", (width, height)) + pixels = image.load() + for i in range(width): + for j in range(height): + pixels[i, j] = ( # pyright: ignore[reportOptionalSubscript] + (i + j * 50) % weights[0], + weights[1], + (i * 300 + j) % weights[2], + ) + return image + + image = create_pattern(width, height, weights) + path = f"/tmp/{extension}-image.{extension}" + image.save(path) + with open(f"/tmp/{extension}-path.json", "w") as f: + json.dump(path, f) + + +@script( + command=["python"], + volume_mounts=[m.VolumeMount(name="tmpdir", mount_path="/tmp")], + outputs=Artifact( + name="hdf5output", + path="/tmp/images.hdf5", + archive=NoneArchiveStrategy(), + ), +) +def to_hdf5(paths: str): + + import h5py # pyright: ignore[reportMissingTypeStubs] + import numpy as np + from PIL import Image + + print("creating hdf5 file") + with h5py.File("/tmp/images.hdf5", "w") as f: + for i, path in enumerate(paths): + path = path.strip('"') + print(f"Got {path}") + with Image.open(path) as image: + arr = np.array(image) + f.create_dataset( # pyright: ignore[reportUnknownMemberType] + f"image_{i}", data=arr, dtype=arr.dtype + ) + print("done") + + +with Workflow( + generate_name="hera-example-in-image-", # name on graphql + entrypoint="workflowentry", + api_version="argoproj.io/v1alpha1", + kind="Workflow", # ClusterWorkflowTemplate", when on graphql + labels={"workflows.diamond.ac.uk/science-group-examples": "true"}, + annotations={ + "workflows.argoproj.io/title": "example remade via hera", + "workflows.argoproj.io/description": """Replicates the functionality of +example.yaml""", + "workflows.diamond.ac.uk/repository": "https://github.com/{% endraw %}{{github_org}}{% raw %}/{% endraw %}{{repo_name}}{% raw %}", + }, + volumes=Volume(name="tmpdir", mount_path="/tmp/", size="1Gi"), +) as w: + with DAG(name="workflowentry"): + params = generate_parameters( + name="params", + arguments={ + "png": "True", + "jpg": "True", + "jpeg": "True", + "tif": "True", + "tiff": "True", + }, + ) + makeimages = create_image(with_param=params.get_parameter("out-parameters")) + makehdf5 = to_hdf5( + arguments={ + "paths": makeimages.get_parameter("out-paths"), + } + ) + params >> makeimages >> makehdf5 # pyright: ignore +with open("example_in_image.txt", "w") as div: + div.write(w.to_yaml()) # pyright: ignore[reportUnknownMemberType] +{% endraw %} diff --git a/src/copier_template/tests/test_keycloak_checker.py.jinja b/src/copier_template/tests/test_keycloak_checker.py.jinja index 2eda176..4b27f07 100644 --- a/src/copier_template/tests/test_keycloak_checker.py.jinja +++ b/src/copier_template/tests/test_keycloak_checker.py.jinja @@ -1,35 +1,120 @@ import os -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch + +from pytest import mark from {{project_name}}.auth.keycloak_checker import set_token_env_variable +@mark.parametrize( + "return_present", + [ + (False), + (True), + ], +) +@patch("{{project_name}}.auth.keycloak_checker.dotenv.load_dotenv") +@patch("{{project_name}}.auth.keycloak_checker.dotenv.set_key") @patch("{{project_name}}.auth.keycloak_checker.KeycloakOpenID") @patch("{{project_name}}.auth.keycloak_checker.generate_code_verifier") @patch("{{project_name}}.auth.keycloak_checker.generate_code_challenge") -@patch("{{project_name}}.auth.keycloak_checker.open_auth_url") -def test_return_key( - mock_open_auth_url: MagicMock, +@patch("{{project_name}}.auth.keycloak_checker.token_expired") +def test_set_token_env_variable( + mock_token_expired: MagicMock, mock_gen_code_challenge: MagicMock, mock_gen_code_verifier: MagicMock, mock_gen_keycloak_id: MagicMock, + mock_set_key: MagicMock, + mock_load_env: MagicMock, + return_present: bool, ): + port=8000 mock_gen_code_verifier.return_value = "verifier" mock_gen_code_challenge.return_value = ("challenge", "S256") os.environ["AUTH"] = "auth_url_code" + os.environ["REFRESHTOKEN"] = "refresh" + keycloak = MagicMock() mock_gen_keycloak_id.return_value = keycloak - + mock_token_expired.return_value = return_present keycloak.auth_url.return_value = "https://mock.site" - keycloak.token.return_value = { + token = { + "access_token": "fake_token", + "refresh_token": "fake_refresh", + } + keycloak.token.return_value = token + keycloak.refresh_token.return_value = token + keycloak.decode_token.return_value = {"exp": 123456789} + assert set_token_env_variable(staging) == "fake_token" + + mock_token_expired.assert_called_once_with("https://mock.site", port) + if return_present: + keycloak.token.assert_called_once_with( + grant_type="authorization_code", + code="auth_url_code", + redirect_uri=f"http://localhost:{port}/", + code_verifier="verifier", + ) + keycloak.refresh_token.assert_not_called() + else: + keycloak.refresh_token.assert_called_once_with("refresh") + keycloak.token.assert_not_called() + mock_set_key.assert_has_calls( + [ + call("src/.env", "EXPIRY", str(123456789 + 1800)), + call("src/.env", "TOKEN", "fake_token"), + call("src/.env", "REFRESHTOKEN", "fake_refresh"), + ] + ) + mock_load_env.assert_called_once_with( + dotenv_path="src/.env", + override=True, + ) + +@patch("{{project_name}}.auth.keycloak_checker.exit") +@patch("{{project_name}}.auth.keycloak_checker.print") +@patch("{{project_name}}.auth.keycloak_checker.dotenv.load_dotenv") +@patch("{{project_name}}.auth.keycloak_checker.dotenv.set_key") +@patch("{{project_name}}.auth.keycloak_checker.KeycloakOpenID") +@patch("{{project_name}}.auth.keycloak_checker.generate_code_verifier") +@patch("{{project_name}}.auth.keycloak_checker.generate_code_challenge") +@patch("{{project_name}}.auth.keycloak_checker.token_expired") +def test_set_token_env_variable_attribute_error( + mock_token_expired: MagicMock, + mock_gen_code_challenge: MagicMock, + mock_gen_code_verifier: MagicMock, + mock_gen_keycloak_id: MagicMock, + mock_set_key: MagicMock, + mock_load_env: MagicMock, + mock_print: MagicMock, + mock_exit: MagicMock, +): + mock_gen_code_verifier.return_value = "verifier" + mock_gen_code_challenge.return_value = ("challenge", "S256") + mock_token_expired.return_value = True + + os.environ["REFRESHTOKEN"] = "refresh" + + keycloak = MagicMock() + mock_gen_keycloak_id.return_value = keycloak + + token = { "access_token": "fake_token", "refresh_token": "fake_refresh", } - assert set_token_env_variable() == "fake_token" - mock_open_auth_url.assert_called_once_with("https://mock.site", 8000) - keycloak.token.assert_called_once_with( - grant_type="authorization_code", - code="auth_url_code", - redirect_uri="http://localhost:8000/", - code_verifier="verifier", + + keycloak.refresh_token.return_value = token + + decoded = MagicMock() + decoded.__getitem__.side_effect = AttributeError + keycloak.decode_token.return_value = decoded + + set_token_env_variable() + + mock_print.assert_called_once_with("ERROR:") + mock_exit.assert_called_once_with(1) + mock_set_key.assert_not_called() + mock_load_env.assert_called_once_with( + dotenv_path="src/.env", + override=True, ) diff --git a/src/copier_template/tests/test_open_auth_url.py.jinja b/src/copier_template/tests/test_open_auth_url.py.jinja index f6d3393..f35aca1 100644 --- a/src/copier_template/tests/test_open_auth_url.py.jinja +++ b/src/copier_template/tests/test_open_auth_url.py.jinja @@ -3,26 +3,42 @@ from unittest.mock import MagicMock, patch from {{project_name}}.auth.open_auth_url import ( CallbackHandler, - open_auth_url, + _open_auth_url, # pyright:ignore + token_expired, ) +@patch("{{project_name}}.auth.open_auth_url.dotenv.set_key") +@patch("{{project_name}}.auth.open_auth_url.time.time") @patch("{{project_name}}.auth.open_auth_url.webbrowser.open") @patch("{{project_name}}.auth.open_auth_url._ReusingHTTPServer") def test_open_auth_url_normal_function( mock_http_server: MagicMock, mock_open_browser: MagicMock, + mock_time: MagicMock, + mock_set_key: MagicMock, ): + mock_time.return_value = 100 + os.environ["EXPIRY"] = "" + server = mock_http_server.return_value server.auth_code = "this_is_your_code" - open_auth_url("url") + + _open_auth_url("url", 5173) server.handle_request.assert_called_once() mock_open_browser.assert_called_once_with("url") + mock_set_key.assert_called_once_with( + "src/.env", + "AUTH", + "this_is_your_code", + ) server.socket.shutdown.assert_called_once() server.server_close.assert_called_once() assert os.environ["AUTH"] == "this_is_your_code" +@patch("{{project_name}}.auth.open_auth_url.dotenv.set_key") +@patch("{{project_name}}.auth.open_auth_url.time.time") @patch("{{project_name}}.auth.open_auth_url.exit") @patch("{{project_name}}.auth.open_auth_url.webbrowser.open") @patch("{{project_name}}.auth.open_auth_url._ReusingHTTPServer") @@ -30,18 +46,64 @@ def test_open_auth_url_raises_error( mock_http_server: MagicMock, mock_open_browser: MagicMock, mock_exit: MagicMock, + mock_time: MagicMock, + mock_set_key: MagicMock, ): + mock_time.return_value = 100 + os.environ["EXPIRY"] = "" + server = mock_http_server.return_value - server.auth_code = "this_is_your_code" server.handle_request.side_effect = OSError - open_auth_url("url") + + _open_auth_url("url", 5173) mock_open_browser.assert_called_once_with("url") assert os.environ["AUTH"] == "" mock_exit.assert_called_once_with(1) + mock_set_key.assert_not_called() server.socket.shutdown.assert_called_once() server.server_close.assert_called_once() +@patch("{{project_name}}.auth.open_auth_url._open_auth_url") +@patch("{{project_name}}.auth.open_auth_url.dotenv.load_dotenv") +@patch("{{project_name}}.auth.open_auth_url.time.time") +def test_token_expired( + mock_time: MagicMock, + mock_load_env: MagicMock, + mock_open_auth_url: MagicMock, +): + mock_time.return_value = 100 + os.environ["EXPIRY"] = "50" + + assert token_expired("url", 5173) is True + + mock_load_env.assert_called_once_with( + dotenv_path="src/.env", + override=True, + ) + mock_open_auth_url.assert_called_once_with("url", 5173) + + +@patch("{{project_name}}.auth.open_auth_url._open_auth_url") +@patch("{{project_name}}.auth.open_auth_url.dotenv.load_dotenv") +@patch("{{project_name}}.auth.open_auth_url.time.time") +def test_token_not_expired( + mock_time: MagicMock, + mock_load_env: MagicMock, + mock_open_auth_url: MagicMock, +): + mock_time.return_value = 100 + os.environ["EXPIRY"] = "200" + + assert token_expired("url", 5173) is False + + mock_load_env.assert_called_once_with( + dotenv_path="src/.env", + override=True, + ) + mock_open_auth_url.assert_not_called() + + def test_handler_normal_function(): handler = CallbackHandler.__new__(CallbackHandler) handler.path = "/?code=this_is_your_code" diff --git a/src/python_interface_to_workflows/auth/keycloak_checker.py b/src/python_interface_to_workflows/auth/keycloak_checker.py index 71390fb..6f6ff11 100644 --- a/src/python_interface_to_workflows/auth/keycloak_checker.py +++ b/src/python_interface_to_workflows/auth/keycloak_checker.py @@ -1,10 +1,20 @@ import os +from typing import TypedDict, cast import dotenv from keycloak import KeycloakOpenID from keycloak.pkce_utils import generate_code_challenge, generate_code_verifier -from python_interface_to_workflows.auth.open_auth_url import open_auth_url +from python_interface_to_workflows.auth.open_auth_url import token_expired + + +class TokenResponse(TypedDict): + access_token: str + refresh_token: str + + +class DecodedToken(TypedDict): + exp: int def set_token_env_variable(staging: bool) -> str: @@ -36,14 +46,38 @@ def set_token_env_variable(staging: bool) -> str: code_challenge=code_challenge, code_challenge_method=code_challenge_method, ) - open_auth_url(auth_url, port) - token: dict[str, str] = ( # pyright: ignore[reportUnknownVariableType] - keycloak_openid.token( # pyright: ignore[reportUnknownMemberType] - grant_type="authorization_code", - code=os.environ["AUTH"], - redirect_uri=f"http://localhost:{port}/", - code_verifier=code_verifier, + if token_expired(auth_url, port): + token = cast( + TokenResponse, + keycloak_openid.token( # pyright: ignore[reportUnknownMemberType] + grant_type="authorization_code", + code=str(os.environ.get("AUTH")), + redirect_uri=f"http://localhost:{port}/", + code_verifier=code_verifier, + ), + ) + else: + token = cast( + TokenResponse, + keycloak_openid.refresh_token( # pyright: ignore[reportUnknownMemberType] + str(os.environ.get("REFRESHTOKEN")) + ), ) + token_info = cast( + DecodedToken, + keycloak_openid.decode_token( # pyright: ignore[reportUnknownMemberType] + token["access_token"] + ), ) - dotenv.set_key("src/.env", "TOKEN", token["access_token"].strip("'")) - return token["access_token"] # pyright: ignore[reportUnknownArgumentType] + try: + expire_time = int(token_info["exp"]) + 1800 + dotenv.set_key("src/.env", "EXPIRY", str(expire_time)) + dotenv.set_key("src/.env", "TOKEN", token["access_token"].strip("'")) + dotenv.set_key("src/.env", "REFRESHTOKEN", token["refresh_token"].strip("'")) + except AttributeError: + print("ERROR:") + exit(1) + finally: + dotenv.load_dotenv(dotenv_path="src/.env", override=True) + + return token["access_token"] diff --git a/src/python_interface_to_workflows/auth/open_auth_url.py b/src/python_interface_to_workflows/auth/open_auth_url.py index 2f0dcf9..f8c2581 100644 --- a/src/python_interface_to_workflows/auth/open_auth_url.py +++ b/src/python_interface_to_workflows/auth/open_auth_url.py @@ -1,5 +1,6 @@ import os import socket +import time import urllib.parse import webbrowser from http.server import BaseHTTPRequestHandler, HTTPServer @@ -28,7 +29,17 @@ def do_GET(self): self.wfile.write(b"Missing authorization code.") -def open_auth_url(auth_url: str, port: int) -> None: +def token_expired(auth_url: str, port: int) -> bool: + dotenv.load_dotenv(dotenv_path="src/.env", override=True) + expiry_str: str = str(os.environ.get("EXPIRY")).strip("'") + if (expiry_str == "" or int(expiry_str)) <= float(time.time()): + _open_auth_url(auth_url, port) + return True + else: + return False + + +def _open_auth_url(auth_url: str, port: int) -> None: httpd = _ReusingHTTPServer(("localhost", port), CallbackHandler) webbrowser.open(auth_url) try: diff --git a/src/python_interface_to_workflows/submit_to_argo.py b/src/python_interface_to_workflows/submit_to_argo.py index 82b86db..585fd97 100644 --- a/src/python_interface_to_workflows/submit_to_argo.py +++ b/src/python_interface_to_workflows/submit_to_argo.py @@ -2,7 +2,6 @@ from typing import cast import dotenv -from hera.shared import global_config from hera.workflows import Workflow, WorkflowsService from hera.workflows import models as m @@ -12,12 +11,11 @@ def submit_workflow_to_argo(w: Workflow): set_token_env_variable(staging=True) dotenv.load_dotenv(dotenv_path="src/.env", override=True) - w.namespace = os.environ.get("NAMESPACE") + w.namespace = os.environ.get("VISIT") w.workflows_service = WorkflowsService( host=str(os.environ.get("HOST")).strip("'"), token=str(os.environ.get("TOKEN")).strip("'"), ) - global_config.image = str(os.environ.get("DEFAULT_IMAGE")) submitted_w = cast(m.Workflow, w.create()) name = submitted_w.metadata.name namespace = submitted_w.metadata.namespace diff --git a/src/python_interface_to_workflows/templates/divisionyaml.yaml b/src/python_interface_to_workflows/templates/divisionyaml.yaml index 34bed16..79b810a 100644 --- a/src/python_interface_to_workflows/templates/divisionyaml.yaml +++ b/src/python_interface_to_workflows/templates/divisionyaml.yaml @@ -1,5 +1,5 @@ apiVersion: argoproj.io/v1alpha1 -kind: Workflow +kind: WorkflowTemplate metadata: generateName: hera-division- annotations: diff --git a/src/python_interface_to_workflows/templates/example.yaml b/src/python_interface_to_workflows/templates/example.txt similarity index 99% rename from src/python_interface_to_workflows/templates/example.yaml rename to src/python_interface_to_workflows/templates/example.txt index 1fbb1b7..857fb41 100644 --- a/src/python_interface_to_workflows/templates/example.yaml +++ b/src/python_interface_to_workflows/templates/example.txt @@ -1,5 +1,5 @@ apiVersion: argoproj.io/v1alpha1 -kind: Workflow +kind: WorkflowTemplate metadata: generateName: hera-example- annotations: diff --git a/src/python_interface_to_workflows/templates/example_in_image.txt b/src/python_interface_to_workflows/templates/example_in_image.txt new file mode 100644 index 0000000..64f9565 --- /dev/null +++ b/src/python_interface_to_workflows/templates/example_in_image.txt @@ -0,0 +1,196 @@ +apiVersion: argoproj.io/v1alpha1 +kind: Workflow +metadata: + generateName: hera-example-in-image- + annotations: + workflows.argoproj.io/description: |- + Replicates the functionality of + example.yaml + workflows.argoproj.io/title: example remade via hera + workflows.diamond.ac.uk/repository: https://github.com/DiamondLightSource/python-interface-to-workflows + labels: + workflows.diamond.ac.uk/science-group-examples: 'true' +spec: + entrypoint: workflowentry + templates: + - name: workflowentry + dag: + tasks: + - name: params + template: generate-parameters + arguments: + parameters: + - name: png + value: 'True' + - name: jpg + value: 'True' + - name: jpeg + value: 'True' + - name: tif + value: 'True' + - name: tiff + value: 'True' + - name: create-image + depends: params + template: create-image + withParam: '{{tasks.params.outputs.parameters.out-parameters}}' + arguments: + parameters: + - name: width + value: '{{item.width}}' + - name: height + value: '{{item.height}}' + - name: weights + value: '{{item.weights}}' + - name: extension + value: '{{item.extension}}' + - name: to-hdf5 + depends: create-image + template: to-hdf5 + arguments: + parameters: + - name: paths + value: '{{tasks.create-image.outputs.parameters.out-paths}}' + - name: generate-parameters + inputs: + parameters: + - name: png + - name: jpg + - name: jpeg + - name: tif + - name: tiff + outputs: + parameters: + - name: out-parameters + valueFrom: + path: /tmp/parameters.json + script: + image: ghcr.io/matt-carre/python-interface-to-workflows-default-image + source: |- + import os + import sys + sys.path.append(os.getcwd()) + import json + try: jpeg = json.loads(r'''{{inputs.parameters.jpeg}}''') + except: jpeg = r'''{{inputs.parameters.jpeg}}''' + try: jpg = json.loads(r'''{{inputs.parameters.jpg}}''') + except: jpg = r'''{{inputs.parameters.jpg}}''' + try: png = json.loads(r'''{{inputs.parameters.png}}''') + except: png = r'''{{inputs.parameters.png}}''' + try: tif = json.loads(r'''{{inputs.parameters.tif}}''') + except: tif = r'''{{inputs.parameters.tif}}''' + try: tiff = json.loads(r'''{{inputs.parameters.tiff}}''') + except: tiff = r'''{{inputs.parameters.tiff}}''' + + import json + params: list[dict[str, int | list[int] | str] | None] = [{'width': 500, 'height': 500, 'weights': [255, 1, 100], 'extension': 'png'} if png.lower() == 'true' else None, {'width': 600, 'height': 200, 'weights': [100, 150, 100], 'extension': 'jpg'} if jpg.lower() == 'true' else None, {'width': 300, 'height': 400, 'weights': [100, 150, 100], 'extension': 'jpeg'} if jpeg.lower() == 'true' else None, {'width': 300, 'height': 200, 'weights': [230, 100, 1], 'extension': 'tif'} if tif.lower() == 'true' else None, {'width': 200, 'height': 300, 'weights': [230, 100, 1], 'extension': 'tiff'} if tiff.lower() == 'true' else None] + params_to_write: list[dict[str, int | list[int] | str]] = [image_params for image_params in params if image_params is not None] + with open('/tmp/parameters.json', 'w') as f: + json.dump(params_to_write, f) + command: + - python + volumeMounts: + - name: tmpdir + mountPath: /tmp + - name: create-image + inputs: + parameters: + - name: width + - name: height + - name: weights + - name: extension + outputs: + artifacts: + - name: '{{inputs.parameters.extension}}-image' + path: /tmp/{{inputs.parameters.extension}}-image.{{inputs.parameters.extension}} + archive: + none: {} + parameters: + - name: out-paths + valueFrom: + path: /tmp/{{inputs.parameters.extension}}-path.json + script: + image: ghcr.io/matt-carre/python-interface-to-workflows-default-image + source: |- + import os + import sys + sys.path.append(os.getcwd()) + import json + try: extension = json.loads(r'''{{inputs.parameters.extension}}''') + except: extension = r'''{{inputs.parameters.extension}}''' + try: height = json.loads(r'''{{inputs.parameters.height}}''') + except: height = r'''{{inputs.parameters.height}}''' + try: weights = json.loads(r'''{{inputs.parameters.weights}}''') + except: weights = r'''{{inputs.parameters.weights}}''' + try: width = json.loads(r'''{{inputs.parameters.width}}''') + except: width = r'''{{inputs.parameters.width}}''' + + import json + from PIL import Image + + def create_pattern(width: int, height: int, weights: tuple[int, int, int]) -> Image.Image: + print(f'width: {width}') + print(f'height: {height}') + print(f'RBG weights: {weights}') + image = Image.new('RGB', (width, height)) + pixels = image.load() + for i in range(width): + for j in range(height): + pixels[i, j] = ((i + j * 50) % weights[0], weights[1], (i * 300 + j) % weights[2]) + return image + image = create_pattern(width, height, weights) + path = f'/tmp/{extension}-image.{extension}' + image.save(path) + with open(f'/tmp/{extension}-path.json', 'w') as f: + json.dump(path, f) + command: + - python + volumeMounts: + - name: tmpdir + mountPath: /tmp + - name: to-hdf5 + inputs: + parameters: + - name: paths + outputs: + artifacts: + - name: hdf5output + path: /tmp/images.hdf5 + archive: + none: {} + script: + image: ghcr.io/matt-carre/python-interface-to-workflows-default-image + source: |- + import os + import sys + sys.path.append(os.getcwd()) + import json + try: paths = json.loads(r'''{{inputs.parameters.paths}}''') + except: paths = r'''{{inputs.parameters.paths}}''' + + import h5py + import numpy as np + from PIL import Image + print('creating hdf5 file') + with h5py.File('/tmp/images.hdf5', 'w') as f: + for i, path in enumerate(paths): + path = path.strip('"') + print(f'Got {path}') + with Image.open(path) as image: + arr = np.array(image) + f.create_dataset(f'image_{i}', data=arr, dtype=arr.dtype) + print('done') + command: + - python + volumeMounts: + - name: tmpdir + mountPath: /tmp + volumeClaimTemplates: + - metadata: + name: tmpdir + spec: + accessModes: + - ReadWriteOnce + resources: + requests: + storage: 1Gi diff --git a/src/python_interface_to_workflows/workflow_definitions/create_division_yaml.py b/src/python_interface_to_workflows/workflow_definitions/create_division_yaml.py index f481093..3f063de 100644 --- a/src/python_interface_to_workflows/workflow_definitions/create_division_yaml.py +++ b/src/python_interface_to_workflows/workflow_definitions/create_division_yaml.py @@ -31,7 +31,7 @@ def do_division(a: int, b: int): generate_name="hera-division-", # when running on graphql this should be name entrypoint="divide", api_version="argoproj.io/v1alpha1", - kind="Workflow", # ClusterWorkflowTemplate", when on graphql + kind="WorkflowTemplate", # ClusterWorkflowTemplate", when on graphql labels={"workflows.diamond.ac.uk/science-group-examples": "true"}, annotations={ "workflows.argoproj.io/title": "Division via hera test", diff --git a/src/python_interface_to_workflows/workflow_definitions/create_example_template.py b/src/python_interface_to_workflows/workflow_definitions/create_example_template.py index e718511..7b9de3c 100644 --- a/src/python_interface_to_workflows/workflow_definitions/create_example_template.py +++ b/src/python_interface_to_workflows/workflow_definitions/create_example_template.py @@ -147,7 +147,7 @@ def to_hdf5(paths: str): generate_name="hera-example-", # when running on graphql this should be name entrypoint="workflowentry", api_version="argoproj.io/v1alpha1", - kind="Workflow", # ClusterWorkflowTemplate", when on graphql + kind="WorkflowTemplate", # ClusterWorkflowTemplate", when on graphql labels={"workflows.diamond.ac.uk/science-group-examples": "true"}, annotations={ "workflows.argoproj.io/title": "example remade via hera", @@ -178,5 +178,5 @@ def to_hdf5(paths: str): [install, params] >> makeimages >> makehdf5 # pyright: ignore -with open("example.yaml", "w") as div: +with open("example.txt", "w") as div: div.write(w.to_yaml()) # pyright: ignore[reportUnknownMemberType] diff --git a/src/python_interface_to_workflows/workflow_definitions/create_example_template_within_default_image.py b/src/python_interface_to_workflows/workflow_definitions/create_example_template_within_default_image.py new file mode 100644 index 0000000..88740ca --- /dev/null +++ b/src/python_interface_to_workflows/workflow_definitions/create_example_template_within_default_image.py @@ -0,0 +1,172 @@ +import os + +from hera.shared import global_config +from hera.workflows import ( + DAG, + Artifact, + Parameter, + Script, + Volume, + Workflow, + script, # pyright: ignore[reportUnknownVariableType] +) +from hera.workflows import models as m +from hera.workflows.archive import NoneArchiveStrategy + +global_config.set_class_defaults( # pyright: ignore + Script, image=str(os.environ.get("DEFAULT_IMAGE")) +) + + +@script( + command=["python"], + outputs=Parameter( + name="out-parameters", value_from=m.ValueFrom(path="/tmp/parameters.json") + ), + volume_mounts=[m.VolumeMount(name="tmpdir", mount_path="/tmp")], +) +def generate_parameters( + png: str, + jpg: str, + jpeg: str, + tif: str, + tiff: str, +): + import json + + params: list[dict[str, int | list[int] | str] | None] = [ + {"width": 500, "height": 500, "weights": [255, 1, 100], "extension": "png"} + if png.lower() == "true" + else None, + {"width": 600, "height": 200, "weights": [100, 150, 100], "extension": "jpg"} + if jpg.lower() == "true" + else None, + {"width": 300, "height": 400, "weights": [100, 150, 100], "extension": "jpeg"} + if jpeg.lower() == "true" + else None, + {"width": 300, "height": 200, "weights": [230, 100, 1], "extension": "tif"} + if tif.lower() == "true" + else None, + {"width": 200, "height": 300, "weights": [230, 100, 1], "extension": "tiff"} + if tiff.lower() == "true" + else None, + ] + params_to_write: list[dict[str, int | list[int] | str]] = [ + image_params for image_params in params if image_params is not None + ] + with open("/tmp/parameters.json", "w") as f: + json.dump(params_to_write, f) + + +@script( + command=["python"], + volume_mounts=[m.VolumeMount(name="tmpdir", mount_path="/tmp")], + outputs=[ + Parameter( + name="out-paths", + value_from=m.ValueFrom( + path="/tmp/{{inputs.parameters.extension}}-path.json" + ), + ), + Artifact( + name="{{inputs.parameters.extension}}-image", + path="/tmp/{{inputs.parameters.extension}}-image.{{inputs.parameters.extension}}", + archive=NoneArchiveStrategy(), + ), + ], +) +def create_image( + width: int, height: int, weights: tuple[int, int, int], extension: str +): + import json + + from PIL import Image + + def create_pattern( + width: int, + height: int, + weights: tuple[int, int, int], + ) -> Image.Image: + print(f"width: {width}") + print(f"height: {height}") + print(f"RBG weights: {weights}") + image = Image.new("RGB", (width, height)) + pixels = image.load() + for i in range(width): + for j in range(height): + pixels[i, j] = ( # pyright: ignore[reportOptionalSubscript] + (i + j * 50) % weights[0], + weights[1], + (i * 300 + j) % weights[2], + ) + return image + + image = create_pattern(width, height, weights) + path = f"/tmp/{extension}-image.{extension}" + image.save(path) + with open(f"/tmp/{extension}-path.json", "w") as f: + json.dump(path, f) + + +@script( + command=["python"], + volume_mounts=[m.VolumeMount(name="tmpdir", mount_path="/tmp")], + outputs=Artifact( + name="hdf5output", + path="/tmp/images.hdf5", + archive=NoneArchiveStrategy(), + ), +) +def to_hdf5(paths: str): + + import h5py # pyright: ignore[reportMissingTypeStubs] + import numpy as np + from PIL import Image + + print("creating hdf5 file") + with h5py.File("/tmp/images.hdf5", "w") as f: + for i, path in enumerate(paths): + path = path.strip('"') + print(f"Got {path}") + with Image.open(path) as image: + arr = np.array(image) + f.create_dataset( # pyright: ignore[reportUnknownMemberType] + f"image_{i}", data=arr, dtype=arr.dtype + ) + print("done") + + +with Workflow( + generate_name="hera-example-in-image-", # name on graphql + entrypoint="workflowentry", + api_version="argoproj.io/v1alpha1", + kind="Workflow", # ClusterWorkflowTemplate", when on graphql + labels={"workflows.diamond.ac.uk/science-group-examples": "true"}, + annotations={ + "workflows.argoproj.io/title": "example remade via hera", + "workflows.argoproj.io/description": """Replicates the functionality of +example.yaml""", + "workflows.diamond.ac.uk/repository": "https://github.com/DiamondLightSource/python-interface-to-workflows", + }, + volumes=Volume(name="tmpdir", mount_path="/tmp/", size="1Gi"), +) as w: + with DAG(name="workflowentry"): + params = generate_parameters( + name="params", + arguments={ + "png": "True", + "jpg": "True", + "jpeg": "True", + "tif": "True", + "tiff": "True", + }, + ) + makeimages = create_image(with_param=params.get_parameter("out-parameters")) + makehdf5 = to_hdf5( + arguments={ + "paths": makeimages.get_parameter("out-paths"), + } + ) + params >> makeimages >> makehdf5 # pyright: ignore +with open("example_in_image.txt", "w") as div: + div.write(w.to_yaml()) # pyright: ignore[reportUnknownMemberType] diff --git a/tests/test_keycloak_checker.py b/tests/test_keycloak_checker.py index bb5c776..c658d0c 100644 --- a/tests/test_keycloak_checker.py +++ b/tests/test_keycloak_checker.py @@ -1,40 +1,124 @@ import os -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch from pytest import mark from python_interface_to_workflows.auth.keycloak_checker import set_token_env_variable -@mark.parametrize("dev,port", [(True, 5173), (False, 8000)]) +@mark.parametrize( + "staging,port,return_present", + [ + (True, 5173, True), + (False, 8000, False), + (True, 5173, False), + (False, 8000, True), + ], +) +@patch("python_interface_to_workflows.auth.keycloak_checker.dotenv.load_dotenv") +@patch("python_interface_to_workflows.auth.keycloak_checker.dotenv.set_key") @patch("python_interface_to_workflows.auth.keycloak_checker.KeycloakOpenID") @patch("python_interface_to_workflows.auth.keycloak_checker.generate_code_verifier") @patch("python_interface_to_workflows.auth.keycloak_checker.generate_code_challenge") -@patch("python_interface_to_workflows.auth.keycloak_checker.open_auth_url") +@patch("python_interface_to_workflows.auth.keycloak_checker.token_expired") def test_set_token_env_variable( - mock_open_auth_url: MagicMock, + mock_token_expired: MagicMock, mock_gen_code_challenge: MagicMock, mock_gen_code_verifier: MagicMock, mock_gen_keycloak_id: MagicMock, - dev: bool, + mock_set_key: MagicMock, + mock_load_env: MagicMock, + staging: bool, port: int, + return_present: bool, ): mock_gen_code_verifier.return_value = "verifier" mock_gen_code_challenge.return_value = ("challenge", "S256") os.environ["AUTH"] = "auth_url_code" + os.environ["REFRESHTOKEN"] = "refresh" + keycloak = MagicMock() mock_gen_keycloak_id.return_value = keycloak - + mock_token_expired.return_value = return_present keycloak.auth_url.return_value = "https://mock.site" - keycloak.token.return_value = { + token = { + "access_token": "fake_token", + "refresh_token": "fake_refresh", + } + keycloak.token.return_value = token + keycloak.refresh_token.return_value = token + keycloak.decode_token.return_value = {"exp": 123456789} + assert set_token_env_variable(staging) == "fake_token" + + mock_token_expired.assert_called_once_with("https://mock.site", port) + if return_present: + keycloak.token.assert_called_once_with( + grant_type="authorization_code", + code="auth_url_code", + redirect_uri=f"http://localhost:{port}/", + code_verifier="verifier", + ) + keycloak.refresh_token.assert_not_called() + else: + keycloak.refresh_token.assert_called_once_with("refresh") + keycloak.token.assert_not_called() + mock_set_key.assert_has_calls( + [ + call("src/.env", "EXPIRY", str(123456789 + 1800)), + call("src/.env", "TOKEN", "fake_token"), + call("src/.env", "REFRESHTOKEN", "fake_refresh"), + ] + ) + mock_load_env.assert_called_once_with( + dotenv_path="src/.env", + override=True, + ) + + +@patch("python_interface_to_workflows.auth.keycloak_checker.exit") +@patch("python_interface_to_workflows.auth.keycloak_checker.print") +@patch("python_interface_to_workflows.auth.keycloak_checker.dotenv.load_dotenv") +@patch("python_interface_to_workflows.auth.keycloak_checker.dotenv.set_key") +@patch("python_interface_to_workflows.auth.keycloak_checker.KeycloakOpenID") +@patch("python_interface_to_workflows.auth.keycloak_checker.generate_code_verifier") +@patch("python_interface_to_workflows.auth.keycloak_checker.generate_code_challenge") +@patch("python_interface_to_workflows.auth.keycloak_checker.token_expired") +def test_set_token_env_variable_attribute_error( + mock_token_expired: MagicMock, + mock_gen_code_challenge: MagicMock, + mock_gen_code_verifier: MagicMock, + mock_gen_keycloak_id: MagicMock, + mock_set_key: MagicMock, + mock_load_env: MagicMock, + mock_print: MagicMock, + mock_exit: MagicMock, +): + mock_gen_code_verifier.return_value = "verifier" + mock_gen_code_challenge.return_value = ("challenge", "S256") + mock_token_expired.return_value = True + + os.environ["REFRESHTOKEN"] = "refresh" + + keycloak = MagicMock() + mock_gen_keycloak_id.return_value = keycloak + + token = { "access_token": "fake_token", "refresh_token": "fake_refresh", } - assert set_token_env_variable(dev) == "fake_token" - mock_open_auth_url.assert_called_once_with("https://mock.site", port) - keycloak.token.assert_called_once_with( - grant_type="authorization_code", - code="auth_url_code", - redirect_uri=f"http://localhost:{port}/", - code_verifier="verifier", + + keycloak.refresh_token.return_value = token + + decoded = MagicMock() + decoded.__getitem__.side_effect = AttributeError + keycloak.decode_token.return_value = decoded + + set_token_env_variable(True) + + mock_print.assert_called_once_with("ERROR:") + mock_exit.assert_called_once_with(1) + mock_set_key.assert_not_called() + mock_load_env.assert_called_once_with( + dotenv_path="src/.env", + override=True, ) diff --git a/tests/test_open_auth_url.py b/tests/test_open_auth_url.py index 59fa0d3..0acd322 100644 --- a/tests/test_open_auth_url.py +++ b/tests/test_open_auth_url.py @@ -3,26 +3,42 @@ from python_interface_to_workflows.auth.open_auth_url import ( CallbackHandler, - open_auth_url, + _open_auth_url, # pyright:ignore + token_expired, ) +@patch("python_interface_to_workflows.auth.open_auth_url.dotenv.set_key") +@patch("python_interface_to_workflows.auth.open_auth_url.time.time") @patch("python_interface_to_workflows.auth.open_auth_url.webbrowser.open") @patch("python_interface_to_workflows.auth.open_auth_url._ReusingHTTPServer") def test_open_auth_url_normal_function( mock_http_server: MagicMock, mock_open_browser: MagicMock, + mock_time: MagicMock, + mock_set_key: MagicMock, ): + mock_time.return_value = 100 + os.environ["EXPIRY"] = "" + server = mock_http_server.return_value server.auth_code = "this_is_your_code" - open_auth_url("url", 5173) + + _open_auth_url("url", 5173) server.handle_request.assert_called_once() mock_open_browser.assert_called_once_with("url") + mock_set_key.assert_called_once_with( + "src/.env", + "AUTH", + "this_is_your_code", + ) server.socket.shutdown.assert_called_once() server.server_close.assert_called_once() assert os.environ["AUTH"] == "this_is_your_code" +@patch("python_interface_to_workflows.auth.open_auth_url.dotenv.set_key") +@patch("python_interface_to_workflows.auth.open_auth_url.time.time") @patch("python_interface_to_workflows.auth.open_auth_url.exit") @patch("python_interface_to_workflows.auth.open_auth_url.webbrowser.open") @patch("python_interface_to_workflows.auth.open_auth_url._ReusingHTTPServer") @@ -30,18 +46,64 @@ def test_open_auth_url_raises_error( mock_http_server: MagicMock, mock_open_browser: MagicMock, mock_exit: MagicMock, + mock_time: MagicMock, + mock_set_key: MagicMock, ): + mock_time.return_value = 100 + os.environ["EXPIRY"] = "" + server = mock_http_server.return_value - server.auth_code = "this_is_your_code" server.handle_request.side_effect = OSError - open_auth_url("url", 5173) + + _open_auth_url("url", 5173) mock_open_browser.assert_called_once_with("url") assert os.environ["AUTH"] == "" mock_exit.assert_called_once_with(1) + mock_set_key.assert_not_called() server.socket.shutdown.assert_called_once() server.server_close.assert_called_once() +@patch("python_interface_to_workflows.auth.open_auth_url._open_auth_url") +@patch("python_interface_to_workflows.auth.open_auth_url.dotenv.load_dotenv") +@patch("python_interface_to_workflows.auth.open_auth_url.time.time") +def test_token_expired( + mock_time: MagicMock, + mock_load_env: MagicMock, + mock_open_auth_url: MagicMock, +): + mock_time.return_value = 100 + os.environ["EXPIRY"] = "50" + + assert token_expired("url", 5173) is True + + mock_load_env.assert_called_once_with( + dotenv_path="src/.env", + override=True, + ) + mock_open_auth_url.assert_called_once_with("url", 5173) + + +@patch("python_interface_to_workflows.auth.open_auth_url._open_auth_url") +@patch("python_interface_to_workflows.auth.open_auth_url.dotenv.load_dotenv") +@patch("python_interface_to_workflows.auth.open_auth_url.time.time") +def test_token_not_expired( + mock_time: MagicMock, + mock_load_env: MagicMock, + mock_open_auth_url: MagicMock, +): + mock_time.return_value = 100 + os.environ["EXPIRY"] = "200" + + assert token_expired("url", 5173) is False + + mock_load_env.assert_called_once_with( + dotenv_path="src/.env", + override=True, + ) + mock_open_auth_url.assert_not_called() + + def test_handler_normal_function(): handler = CallbackHandler.__new__(CallbackHandler) handler.path = "/?code=this_is_your_code"