Added integration tests against simulation workflow

Author: Michael Smit

.github/workflows/gcp-deploy.reusable.yml

@@ -95,6 +95,11 @@ jobs:
         python-version: "3.11"
     - name: Set up poetry
       run: uv pip install poetry --system
+    - name: Authenticate as tester SA in GCP
+      uses: "google-github-actions/auth@v2"
+      with:
+        workload_identity_provider: "${{ vars._GITHUB_IDENTITY_POOL_PROVIDER_NAME }}"
+        service_account: "tester@${{ vars.PROJECT_ID }}.iam.gserviceaccount.com"
     - name: Auth as tester SA in GCP for full api
       id: get-full-id-token
       uses: "google-github-actions/auth@v2"

common.mk

@@ -1,3 +1,5 @@
+-include ../../terraform/.bootstrap_settings/project.env
+
 build: remove_artifacts install checkformat pyright generate test 
 
 remove_artifacts:

projects/policyengine-apis-integ/Makefile

@@ -1,4 +1,7 @@
 include ../../common.mk
+US_COUNTRY_PACKAGE_VERSION := $(shell ../../projects/policyengine-api-simulation/dump_package_version.sh policyengine-us)
+UK_COUNTRY_PACKAGE_VERSION := $(shell ../../projects/policyengine-api-simulation/dump_package_version.sh policyengine-uk)
+
 test: integ-test
 	echo "Skipping default test target in favor of integ-test"
 
@@ -8,4 +11,7 @@ integ-test:
 	$(if $(SIMULATION_API_ACCESS_TOKEN),SIMULATION_INTEG_TEST_ACCESS_TOKEN='$(SIMULATION_API_ACCESS_TOKEN)') \
 	$(if $(FULL_API_URL),FULL_INTEG_TEST_BASE_URL='$(FULL_API_URL)') \
 	$(if $(SIMULATION_API_URL),SIMULATION_INTEG_TEST_BASE_URL='$(SIMULATION_API_URL)') \
+	$(if $(TF_VAR_project_id),WORKFLOW_INTEG_TEST_PROJECT_ID='$(TF_VAR_project_id)') \
+	$(if $(PROJECT_ID),WORKFLOW_INTEG_TEST_PROJECT_ID='$(PROJECT_ID)') \
+        WORKFLOW_INTEG_TEST_US_MODEL_VERSION='$(US_COUNTRY_PACKAGE_VERSION)' \
 	uv run pytest -n auto

projects/policyengine-apis-integ/pyproject.toml

@@ -8,6 +8,7 @@ dependencies = [
     "policyengine_full_api_client",
     "policyengine_simulation_api_client",
     "pydantic-settings (>=2.8.1,<3.0.0)",
+    "google-cloud-workflows>=1.18.1",
 ]
 
 [project.optional-dependencies]

projects/policyengine-apis-integ/tests/simulation_workflow/__init__.py

@@ -0,0 +1,27 @@
+from pydantic_settings import BaseSettings, SettingsConfigDict
+import pytest
+from .simplified_workflow_client import SimplifiedWorkflowClient
+
+
+class Settings(BaseSettings):
+    project_id: str = "UNKNOWN_PROJECT_ID"
+    location: str = "us-central1"
+    workflow_id: str = "simulation-workflow"
+    us_model_version: str = "UNKNOWN_US_MODEL_VERSION"
+
+    model_config = SettingsConfigDict(env_prefix="workflow_integ_test_")
+
+
+settings = Settings()
+
+
+@pytest.fixture()
+def us_model_version() -> str:
+    return settings.us_model_version
+
+
+@pytest.fixture()
+def client() -> SimplifiedWorkflowClient:
+    return SimplifiedWorkflowClient(
+        settings.project_id, settings.location, settings.workflow_id
+    )

projects/policyengine-apis-integ/tests/simulation_workflow/conftest.py

@@ -0,0 +1,34 @@
+import json
+from typing import Any
+
+from google.cloud import workflows_v1
+from google.cloud.workflows import executions_v1
+
+from google.cloud.workflows.executions_v1.types import executions
+
+
+class SimplifiedWorkflowClient:
+    def __init__(self, project_id: str, location: str, workflow_id: str):
+        self.execution_client = executions_v1.ExecutionsClient()
+        self.workflows_client = workflows_v1.WorkflowsClient()
+        self.project_id = project_id
+        self.location = location
+        self.workflow_id = workflow_id
+
+    def execute(self, argument: Any) -> executions.Execution:
+        parent = self.workflows_client.workflow_path(
+            self.project_id, self.location, self.workflow_id
+        )
+        response = self.execution_client.create_execution(
+            parent=parent,
+            execution=executions.Execution(argument=json.dumps(argument)),
+        )
+
+        execution_finished = False
+        while True:
+            execution = self.execution_client.get_execution(
+                request={"name": response.name}
+            )
+            execution_finished = execution.state != executions.Execution.State.ACTIVE
+            if execution_finished:
+                return execution

projects/policyengine-apis-integ/tests/simulation_workflow/simplified_workflow_client.py

@@ -0,0 +1,39 @@
+from .simplified_workflow_client import SimplifiedWorkflowClient
+from google.cloud.workflows.executions_v1.types import executions
+
+
+def test_calculate_default_model(client: SimplifiedWorkflowClient):
+    execution = client.execute(
+        argument={
+            "country": "us",  # don't use uk. It will try to load extra stuff from huggingface
+            "scope": "macro",
+            "reform": {
+                "gov.irs.credits.ctc.refundable.fully_refundable": {
+                    "2023-01-01.2100-12-31": True
+                }
+            },
+            "subsample": 200,  # reduce the number of households to speed things up.
+            "data": "gs://policyengine-us-data/cps_2023.h5",  # force the service to use google storage (policyengine.py defaults to huggingface)
+        }
+    )
+    assert execution.state == executions.Execution.State.SUCCEEDED
+
+
+def test_calculate_specific_model(
+    client: SimplifiedWorkflowClient, us_model_version: str
+):
+    execution = client.execute(
+        argument={
+            "country": "us",  # don't use uk. It will try to load extra stuff from huggingface
+            "model_version": us_model_version,
+            "scope": "macro",
+            "reform": {
+                "gov.irs.credits.ctc.refundable.fully_refundable": {
+                    "2023-01-01.2100-12-31": True
+                }
+            },
+            "subsample": 200,  # reduce the number of households to speed things up.
+            "data": "gs://policyengine-us-data/cps_2023.h5",  # force the service to use google storage (policyengine.py defaults to huggingface)
+        }
+    )
+    assert execution.state == executions.Execution.State.SUCCEEDED