Source code for operators.dbt.dbt_test

# dbt_test.py
from airflow.utils.decorators import apply_defaults
from airflow.utils.context import Context

from operators.dbt.dbt_base import BaseDBTOperator




class DBTTestOperator(BaseDBTOperator):
    """
    Execute `dbt test` on the remote server.

    :param models: list of model names (without .sql) to test; defaults to all
    :param vars:   dict of vars to pass into dbt
    """

    ui_color = "#DB0058"
    ui_fgcolor = "#FFFFFF"
    custom_operator_name = "DBTTestOperator"

    @apply_defaults
    def __init__(
        self,
        **kwargs,
    ):
        super().__init__(**kwargs)



    def execute(self, context: Context):
        import time

        super().execute(context)

        # build and run
        cmd = " ".join(self.build_command("dbt test"))
        self.log.info(f"Running: {cmd}")

        clock_start = time.time()
        try:
            # When we create a new table, DBT doesn't get notified automatically.
            # So we try to run the command, and if compilation fails, we refresh the source tables and try again.
            # We don't refresh sources every time because it adds ~45 seconds to the run time.
            try:
                self.run_ssh_command(cmd)
            except Exception as e:
                print(e)
                # If less than a minute has passed since we started, try refreshing sources and re-running
                # If its been more than a minute, then it isn't a compilation issue, so just fail out
                if time.time() - clock_start < 60:
                    self.log.info(
                        "dbt test failed, attempting to refresh sources and re-run..."
                    )
                    super().run_dbt_refresh_sources()
                    self.run_ssh_command(cmd)
                else:
                    raise
        except Exception:
            self.log.error("dbt test failed!")
            raise
        finally:
            import json, base64

            test_results = {"passed": [], "failed": [], "warned": [], "errored": []}
            try:
                run_results_output = self.run_ssh_command(
                    f"cat {self.target_path}/run_results.json | base64 -w 0"
                )
                decoded = base64.b64decode(run_results_output).decode("utf-8")
                run_results = json.loads(decoded)

                for result in run_results.get("results", []):
                    test_name = result.get("unique_id", "unknown")
                    status = result.get("status", "unknown")
                    entry = {
                        "test_name": test_name,
                        "status": status,
                        "failures": result.get("failures", 0),
                        "message": result.get("message", ""),
                    }
                    if status == "pass":
                        test_results["passed"].append(entry)
                    elif status == "fail":
                        test_results["failed"].append(entry)
                    elif status == "warn":
                        test_results["warned"].append(entry)
                    else:
                        test_results["errored"].append(entry)

                self.log.info(
                    f"Test results: {len(test_results['passed'])} passed, "
                    f"{len(test_results['failed'])} failed, "
                    f"{len(test_results['warned'])} warned, "
                    f"{len(test_results['errored'])} errored"
                )
            except Exception as e:
                self.log.warning(f"Could not parse test results: {e}")

            result = {
                "job_id": self.job_id,
                "target_path": self.target_path,
                "test_results": test_results,
            }
            context["ti"].xcom_push(key="return_value", value=result)

        self.log.info("dbt test completed successfully.")
        return result