import os
import pytest
from unittest.mock import patch
from aare.common.beamline import mx_beamline, MXBeamline

def test_mx_beamline_default():
    with patch.dict(os.environ, {}, clear=True):
        # If BEAMLINE is not set, it should return SIMULATED
        assert mx_beamline() == MXBeamline.SIMULATED

def test_mx_beamline_x10sa():
    with patch.dict(os.environ, {"BEAMLINE": "x10sa"}):
        assert mx_beamline() == MXBeamline.X10SA

def test_mx_beamline_x06sa():
    with patch.dict(os.environ, {"BEAMLINE": "X06SA "}):
        assert mx_beamline() == MXBeamline.X06SA

def test_mx_beamline_invalid():
    with patch.dict(os.environ, {"BEAMLINE": "INVALID"}):
        assert mx_beamline() == MXBeamline.SIMULATED