Initialization of test profile

CONTRIBUTING.md

@@ -1,25 +1,51 @@
-# Runinng tests
+# Run tests
 
-## Testing environment
+## Quick start
+
+To run only unit tests that do not require and Airflow services to be started.
+
+```bash
+# 1. Start PostgreSQL with docker compose
+hatch run hatch-test.py3.11:start-psql-service
+
+# 2. Create AiiDA profile and databases
+hatch run hatch-test.py3.11:setup-profile
+
+# 4. Run unit tests
+hatch run hatch-test.py3.11:unit-tests
+```
+
+## Unit tests
+
+### Setup test environment
 
 The test environment (`hatch-test`) is configured in `pyproject.toml` with:
 - PostgreSQL connection on port 5434
-- Airflow home directory in `.pytest/airflow`
+- AiiDA configuration in `.pytest/.aiida/`
+- AiiDA profile named `test`
+- Airflow home directory in `.pytest/.aiida/test/airflow/`
 - Custom DAG bundle configuration
 - Test dependencies (pytest, psycopg2-binary, asyncpg)
 
-## Start PostgreSQL service
+The `hatch-test` environment automatically sets these required environment variables.
+Please check the environment variables in the `hatch-test` environment and how the scripts use them if you want to adapt them to your own.
+**Important:** Tests and setup scripts **require** these environment variables. They are automatically set when using the scripts `hatch run hatch-test.py3.11:*` commands.
+If you want to run tests manually outside of hatch, you must set them yourself.
+
+#### 1. Start PostgreSQL service
 
-The test environment uses PostgreSQL on port 5434. Start the PostgreSQL service:
+The test environment uses PostgreSQL. Start the PostgreSQL service:
 
 ```bash
-docker compose -f docker-compose.test.yml up -d
+hatch run hatch-test.py3.11:start-psql-service
 ```
 
+This will start PostgreSQL on port 5434 with the admin user `postgres` (password: `postgres`).
+
 Check if the service is running:
 
 ```bash
-docker ps
+hatch run hatch-test.py3.11:status-psql-service
 ```
 
 Expected output:
@@ -28,28 +54,68 @@ CONTAINER ID   IMAGE         COMMAND                  CREATED      STATUS
 6de79ae7f116   postgres:14   "docker-entrypoint.s…"   2 days ago   Up 2 days (healthy)   0.0.0.0:5434->5432/tcp, [::]:5434->5432/tcp   airflow-provider-aiida-postgres-1
 ```
 
-## Run unit tests
+#### 2. Setup AiiDA test profile
+
+Create the AiiDA test profile that uses the PostgreSQL on port 5434:
+
+```bash
+hatch run hatch-test.py3.11:setup-profile
+```
+
+This script will:
+1. Create a PostgreSQL user `aiida-test` with password `password`
+2. Create two databases owned by `aiida-test`:
+   - `aiida-test`: AiiDA database
+   - `airflow-test`: Airflow database
+3. Create an AiiDA profile named `test` using the `aiida-test` database
+4. Initialize the AiiDA storage (create tables)
+5. Set the profile as the default AiiDA profile
+6. Initialize the Airflow database (run migrations)
+7. Set up the Airflow home directory structure
+
+**Database Hierarchy:**
+```
+postgres (admin)
+  └── aiida-test (user created by setup-profile)
+       ├── aiida-test (database for AiiDA)
+       └── airflow-test (database for Airflow)
+```
+
+The profile will be located at `.pytest/.aiida/test/` (or `$AIIDA_PATH/.aiida/test/` if `AIIDA_PATH` is set) with the following structure:
+```
+.pytest/.aiida/test/
+└── airflow/                    # Airflow files
+    ├── dags/                   # DAG files
+    └── airflow.cfg             # Airflow configuration
+```
+
+**Note:** If the profile already exists, the script will recreate it. If you choose to recreate, it will:
+- Remove the existing AiiDA profile configuration
+- Drop both aiida and airflow databases
+- Drop the aiida `test` user
+- Recreate everything from scratch
+
+## Run unit test
 
 For regular unit tests that don't require Airflow services:
 
 ```bash
-hatch test
+hatch run hatch-test.py3.11:unit-tests
 ```
 
-This runs all non-integration tests using pytest.
-To clean test artifacts including airflow.cfg and logs.
+Or use hatch test with pytest args:
+
 ```bash
-hatch run hatch-test.py3.11:clean
+hatch test -- -m 'not integration'
 ```
 
-## Integration tests
+## Run integration tests
 
 Integration tests require running Airflow services (scheduler, triggerer, dag-processor). These tests are marked with `@pytest.mark.integration`.
 
+### Start Airflow services in foreground (recommended for debugging)
 
-### Start airflow services
-
-In separate terminal windows/tabs, start each service:
+If you prefer to start services manually in separate terminals:
 
 ```bash
 # Terminal 1: Scheduler
@@ -67,6 +133,8 @@ hatch run hatch-test.py3.11:api-server
 
 ### Run integration tests
 
+Once services are running, run the integration tests:
+
 ```bash
 hatch run hatch-test.py3.11:integration
 ```
@@ -76,3 +144,24 @@ Or use hatch test with pytest args:
 ```bash
 hatch test -- -m integration
 ```
+
+##  Run all tests
+
+After setup you can run all tests using
+
+```bash
+hatch test
+```
+
+## Clean test artifacts
+
+Be sure that all airflow services and the docker service have been stopped.
+The docker service can be stopped with.
+```
+hatch run hatch-test.py3.11:stop-psql-service
+```
+
+To clean test artifacts including the PostgreSQL databasese cluster, as well as the aiida and the airflow config.
+```bash
+hatch run hatch-test.py3.11:clean
+```

docker-compose.test.yml

@@ -7,15 +7,15 @@ services:
   postgres:
     image: postgres:14
     environment:
-      POSTGRES_USER: airflow
-      POSTGRES_PASSWORD: airflow
-      POSTGRES_DB: airflow
+      POSTGRES_USER: postgres
+      POSTGRES_PASSWORD: postgres
+      POSTGRES_DB: postgres
     ports:
-      - "${AIRFLOW_PROVIDER_AIIDA__POSTGRES_HOST_PORT:-5432}:5432"
+      - "${POSTGRES_HOST_PORT}:5432"
     volumes:
       - postgres-db-volume:/var/lib/postgresql/data
     healthcheck:
-      test: ["CMD", "pg_isready", "-U", "airflow"]
+      test: ["CMD", "pg_isready", "-U", "postgres"]
       interval: 10s
       retries: 5
       start_period: 5s

examples/trigger_arithmetic_add.py

@@ -1,7 +1,7 @@
 # Create Process
 from airflow_provider_aiida.aiida_core.engine.launch import submit
 from aiida.calculations.arithmetic.add import ArithmeticAddCalculation
-from aiida import load_profile
+from airflow_provider_aiida.aiida_core import load_profile
 load_profile()
 from aiida.orm import load_code, Int
 code = load_code('bash@localhost')

examples/trigger_multiply_add.py

@@ -1,7 +1,7 @@
 # Create Process
 from airflow_provider_aiida.aiida_core.engine.launch import submit
 from aiida.workflows.arithmetic.multiply_add import MultiplyAddWorkChain 
-from aiida import load_profile
+from airflow_provider_aiida.aiida_core import load_profile
 load_profile()
 from aiida.orm import load_code, Int
 code = load_code('bash@localhost')

pyproject.toml

@@ -53,7 +53,7 @@ allow-direct-references = true
 python = "3.11"
 
 [tool.hatch.envs.hatch-test]
-default-args = ["tests"]
+default-args = []
 dependencies = [
     "pytest>=7.0",
     "pytest-mock>=3.10",
@@ -66,28 +66,54 @@ dependencies = [
 python = ["3.11"]
 
 [tool.hatch.envs.hatch-test.env-vars]
-AIRFLOW_HOME = "{root}/.pytest/airflow"
-AIRFLOW_PROVIDER_AIIDA__POSTGRES_HOST_PORT = "5434"
-AIRFLOW__DATABASE__SQL_ALCHEMY_CONN = "postgresql+psycopg2://airflow:airflow@127.0.0.1:5434/airflow"
-AIRFLOW__CORE__DAGS_FOLDER = "{root}/src/airflow_provider_aiida/example_dags"
-AIRFLOW__CORE__LOAD_EXAMPLES = "False"
-AIRFLOW__CORE__LOAD_DEFAULT_CONNECTIONS = "False"
-AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST = '[{{"name":"dags-folder","classpath":"airflow.dag_processing.bundles.local.LocalDagBundle","kwargs":{{}}}},{{"name":"aiida_dags","classpath":"airflow_provider_aiida.bundles.aiida_dag_bundle.AiidaDagBundle","kwargs":{{}}}}]'
+# PostgreSQL configuration (shared by Airflow and AiiDA)
+POSTGRES_HOST = "127.0.0.1"
+POSTGRES_HOST_PORT = "5434"
+POSTGRES_USER = "postgres"
+POSTGRES_PASSWORD = "postgres"
+
+
+# AiiDA configuration
+AIIDA_PATH = "{root}/.pytest"
+AIRFLOW_PROVIDER_AIIDA__TESTS__AIIDA_PROFILE = "test"
+AIRFLOW_PROVIDER_AIIDA__TESTS__AIRFLOW_POSTGRES_PASSWORD = "password"
 
 [tool.hatch.envs.hatch-test.scripts]
-# Start Airflow services (for integration testing)
-scheduler = "airflow scheduler"
-api-server = "airflow api-server"
-triggerer = "airflow triggerer"
-dag-processor = "airflow dag-processor"
+# Setup test environment
+setup-profile = [
+  "python scripts/cmd_profile_create.py --postgres-host {env:POSTGRES_HOST} --postgres-port {env:POSTGRES_HOST_PORT} --postgres-user {env:POSTGRES_USER} --postgres-password {env:POSTGRES_PASSWORD} --profile-name {env:AIRFLOW_PROVIDER_AIIDA__TESTS__AIIDA_PROFILE} --aiida-password {env:AIRFLOW_PROVIDER_AIIDA__TESTS__AIRFLOW_POSTGRES_PASSWORD}",
+]
+reserialize = [
+  "python scripts/cmd_airflow.py --profile-name {env:AIRFLOW_PROVIDER_AIIDA__TESTS__AIIDA_PROFILE} dags reserialize --bundle-name aiida_dags",
+]
+
+teardown-profile = [
+  "python scripts/cmd_profile_delete.py --profile-name {env:AIRFLOW_PROVIDER_AIIDA__TESTS__AIIDA_PROFILE} --postgres-user {env:POSTGRES_USER} --postgres-password {env:POSTGRES_PASSWORD}"
+]
+
+# Process manager commands
+start-psql-service = "docker compose -f docker-compose.test.yml up -d"
+stop-psql-service = "docker compose -f docker-compose.test.yml down -v"
+status-psql-service = "docker ps"
+
+# Start individual Airflow services (for manual testing)
+scheduler = "python scripts/cmd_airflow.py --profile-name {env:AIRFLOW_PROVIDER_AIIDA__TESTS__AIIDA_PROFILE} scheduler"
+api-server = "python scripts/cmd_airflow.py --profile-name {env:AIRFLOW_PROVIDER_AIIDA__TESTS__AIIDA_PROFILE} api-server"
+triggerer = "python scripts/cmd_airflow.py --profile-name {env:AIRFLOW_PROVIDER_AIIDA__TESTS__AIIDA_PROFILE} triggerer"
+dag-processor = "python scripts/cmd_airflow.py --profile-name {env:AIRFLOW_PROVIDER_AIIDA__TESTS__AIIDA_PROFILE} dag-processor"
 
 # Run integration tests
-integration = "pytest -m integration tests"
+unit-tests = "pytest -m 'not integration' {args:{root}/tests}"
+integration-tests = "pytest -m 'integration' {args:{root}/tests}"
+# Default test runner - excludes integration tests, only runs tests/ directory
+run = "pytest {args:{root}/tests}"
 
 # Clean test artifacts
+# NOTE: Only use after services have been stopped to
 clean = [
     "rm -rf .pytest",
-    "echo 'Cleaned .pytest directory'",
+    "docker volume rm -f airflow-provider-aiida_postgres-db-volume",
+    "echo 'Cleaned Docker volumes and .pytest directory'",
 ]
 
 [tool.setuptools]
@@ -103,6 +129,9 @@ exclude = [
 ]
 
 [tool.pytest.ini_options]
+markers = [
+    "integration: marks tests that require Airflow services to be running (scheduler, triggerer)",
+]
 filterwarnings = [
     # AiiDA SQLAlchemy warnings - these are expected and safe to ignore as done in aiida-core
     "ignore:Object of type .* not in session, .* operation along .* will not proceed:sqlalchemy.exc.SAWarning",

scripts/cmd_airflow.py

@@ -0,0 +1,84 @@
+#!/usr/bin/env python
+"""
+Run Airflow CLI commands using an AiiDA profile's configuration.
+
+This script loads an AiiDA profile (default or specified), extracts the Airflow environment
+configuration, and executes the provided Airflow CLI command with those settings.
+
+Example usage:
+    python scripts/cmd_airflow.py db migrate
+    python scripts/cmd_airflow.py dags list
+    python scripts/cmd_airflow.py scheduler
+    python scripts/cmd_airflow.py --profile-name test scheduler
+"""
+
+import sys
+import os
+import subprocess
+from airflow_provider_aiida.aiida_core import load_profile
+
+
+def main():
+    """Main function."""
+    # Extract --profile-name argument if present, keeping other args for Airflow
+    profile_name = None
+    airflow_args = []
+
+    i = 0
+    args = sys.argv[1:]
+    while i < len(args):
+        if args[i] == '--profile-name':
+            if i + 1 < len(args):
+                profile_name = args[i + 1]
+                i += 2  # Skip both --profile-name and its value
+            else:
+                print("Error: --profile-name requires a value")
+                return 1
+        else:
+            airflow_args.append(args[i])
+            i += 1
+
+    if not airflow_args:
+        print("Usage: {} [--profile-name PROFILE] [airflow_command] [airflow_args...]".format(sys.argv[0]))
+        print("\nOptions:")
+        print("  --profile-name PROFILE    Use specified AiiDA profile (default: default profile)")
+        print("\nExamples:")
+        print("  {} db migrate                      # Initialize/migrate Airflow database".format(sys.argv[0]))
+        print("  {} dags list                       # List all DAGs".format(sys.argv[0]))
+        print("  {} dags reserialize -B aiida_dags  # Reserialize DAGs".format(sys.argv[0]))
+        print("  {} scheduler                       # Start the scheduler".format(sys.argv[0]))
+        print("  {} triggerer                       # Start the triggerer".format(sys.argv[0]))
+        print("  {} --profile-name test scheduler   # Start scheduler with 'test' profile".format(sys.argv[0]))
+        return 1
+
+    try:
+        # Load the AiiDA profile
+        from airflow_provider_aiida.aiida_core import load_profile
+        load_profile(profile_name)
+
+
+        # Build the airflow command
+        airflow_command = ['airflow'] + airflow_args
+
+        print(f"Running: {' '.join(airflow_command)}")
+        print("=" * 60)
+        print()
+
+        # Execute the airflow command with the configured environment
+        result = subprocess.run(
+            airflow_command,
+            env=os.environ
+        )
+
+        # Return the same exit code as the airflow command
+        return result.returncode
+
+    except Exception as e:
+        print(f"\n✗ Error: {e}", file=sys.stderr)
+        import traceback
+        traceback.print_exc()
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())