import unittest

from datahub.utilities.urns.data_flow_urn import DataFlowUrn
from datahub.utilities.urns.error import InvalidUrnError


class TestDataFlowUrn(unittest.TestCase):
    def test_parse_urn(self) -> None:
        data_flow_urn_str = "urn:li:dataFlow:(airflow,def,prod)"
        data_flow_urn = DataFlowUrn.create_from_string(data_flow_urn_str)
        assert data_flow_urn.get_orchestrator_name() == "airflow"
        assert data_flow_urn.get_flow_id() == "def"
        assert data_flow_urn.get_env() == "prod"
        assert data_flow_urn.__str__() == "urn:li:dataFlow:(airflow,def,prod)"
        assert data_flow_urn == DataFlowUrn("dataFlow", ["airflow", "def", "prod"])

    def test_invalid_urn(self) -> None:
        with self.assertRaises(InvalidUrnError):
            DataFlowUrn.create_from_string("urn:li:abc:(airflow,def,prod)")

        with self.assertRaises(InvalidUrnError):
            DataFlowUrn.create_from_string("urn:li:dataFlow:(airflow,flow_id)")

        with self.assertRaises(InvalidUrnError):
            DataFlowUrn.create_from_string(
                "urn:li:dataFlow:(airflow,flow_id,invalidEnv)"
            )