add integration tests for pyflyte run and register

Signed-off-by: redartera <[email protected]>

Author: redartera

tests/flytekit/integration/remote/test_remote.py

@@ -17,12 +17,15 @@
 import pytest
 from unittest import mock
 from dataclasses import dataclass
+import random
+import string
 
 from flytekit import LaunchPlan, kwtypes, WorkflowExecutionPhase, task, workflow
 from flytekit.configuration import Config, ImageConfig, SerializationSettings
 from flytekit.core.launch_plan import reference_launch_plan
 from flytekit.core.task import reference_task
 from flytekit.core.workflow import reference_workflow
+from flytekit.models import task as task_models
 from flytekit.exceptions.user import FlyteAssertion, FlyteEntityNotExistException
 from flytekit.extras.sqlite3.task import SQLite3Config, SQLite3Task
 from flytekit.remote.remote import FlyteRemote
@@ -1170,3 +1173,98 @@ def test_register_wf_twice(register):
         ]
     )
     assert out.returncode == 0
+
+
+def test_register_wf_with_default_resources_override(register):
+    # Save the version here to retrieve the created task later
+    version = str(uuid.uuid4())
+    # Register the workflow with overridden default resources
+    out = subprocess.run(
+        [
+            "pyflyte",
+            "--verbose",
+            "-c",
+            CONFIG,
+            "register",
+            "--default-resources",
+            "cpu=1300m;mem=1100Mi",
+            "--image",
+            IMAGE,
+            "--project",
+            PROJECT,
+            "--domain",
+            DOMAIN,
+            "--version",
+            version,
+            MODULE_PATH / "hello_world.py",
+        ]
+    )
+    assert out.returncode == 0
+
+    # Retrieve the created task
+    remote = FlyteRemote(Config.auto(config_file=CONFIG), PROJECT, DOMAIN)
+    task = remote.fetch_task(name="basic.hello_world.say_hello", version=version)
+    assert task.template.container is not None
+    assert task.template.container.resources == task_models.Resources(
+        requests=[
+            task_models.Resources.ResourceEntry(
+                name=task_models.Resources.ResourceName.CPU,
+                value="1300m",
+            ),
+            task_models.Resources.ResourceEntry(
+                name=task_models.Resources.ResourceName.MEMORY,
+                value="1100Mi",
+            ),
+        ],
+        limits=[],
+    )
+
+
+def test_run_wf_with_default_resources_override(register):
+    # Save the execution id here to retrieve the created execution later
+    prefix = random.choice(string.ascii_lowercase)
+    short_random_part = uuid.uuid4().hex[:8]
+    execution_id = f"{prefix}{short_random_part}"
+    # Register the workflow with overridden default resources
+    out = subprocess.run(
+        [
+            "pyflyte",
+            "--verbose",
+            "-c",
+            CONFIG,
+            "run",
+            "--remote",
+            "--default-resources",
+            "cpu=500m;mem=1Gi",
+            "--project",
+            PROJECT,
+            "--domain",
+            DOMAIN,
+            "--name",
+            execution_id,
+            MODULE_PATH / "hello_world.py",
+            "my_wf"
+        ]
+    )
+    assert out.returncode == 0
+
+    # Retrieve the created task
+    remote = FlyteRemote(Config.auto(config_file=CONFIG), PROJECT, DOMAIN)
+    execution = remote.fetch_execution(name=execution_id)
+    execution = remote.wait(execution=execution)
+    version = execution.spec.launch_plan.version
+    task = remote.fetch_task(name="basic.hello_world.say_hello", version=version)
+    assert task.template.container is not None
+    assert task.template.container.resources == task_models.Resources(
+        requests=[
+            task_models.Resources.ResourceEntry(
+                name=task_models.Resources.ResourceName.CPU,
+                value="500m",
+            ),
+            task_models.Resources.ResourceEntry(
+                name=task_models.Resources.ResourceName.MEMORY,
+                value="1Gi",
+            ),
+        ],
+        limits=[],
+    )