import boto3
import pytest
import os
from metaflow import Runner

def test_batch_tag_failures():
    try:
        failed = True
        run = Runner(flow_file=os.path.join(os.path.dirname(__file__), "BatchTags.py"), decospecs=["batch"]).run()
        failed = False
    except Exception:
        pass

    if not failed:
        raise Exception("run should have failed with invalid tags")

def test_batch_tagging_success():
    run = Runner(flow_file=os.path.join(os.path.dirname(__file__), "BatchTags2.py"), decospecs=["batch"]).run()

    assert run.status == "successful"

    start_batch_id = run["start"].task.metadata_dict.get("aws-batch-job-id")
    hello_batch_id = run["hello"].task.metadata_dict.get("aws-batch-job-id")
    end_batch_id = run["end"].task.metadata_dict.get("aws-batch-job-id")

    start_tags = get_aws_tags_for_job(start_batch_id)

    hello_tags = get_aws_tags_for_job(hello_batch_id)

    end_tags = get_aws_tags_for_job(end_batch_id)

    # Add asserts for expected tags





def get_aws_tags_for_job(job_id):
    client = boto3.client("batch")
    resp = client.describe_jobs(jobs=[job_id])
    try:
        return resp["jobs"][0]["tags"]
    except IndexError:
        return {}