import os
import sys
import vetiver

# Load data, model
X_df, y = vetiver.get_mock_data()
model = vetiver.get_mock_model().fit(X_df, y)


def test_vetiver_write_app_no_rspm():
    py_version = str(sys.version_info.major) + "." + str(sys.version_info.minor)
    file = "Dockerfile"
    vetiver.write_docker(rspm_env=False, path="./")
    contents = open(file).read()
    os.remove(file)
    assert (
        contents
        == f"""# # Generated by the vetiver package; edit with care
# start with python base image
FROM python:{py_version}

# create directory in container for vetiver files
WORKDIR /vetiver

# copy  and install requirements
COPY vetiver_requirements.txt /vetiver/requirements.txt

#
RUN pip install --no-cache-dir --upgrade -r /vetiver/requirements.txt

# copy app file
COPY app.py /vetiver/app/app.py

# expose port
EXPOSE 8080

# run vetiver API
CMD ["uvicorn", "app.app:api", "--host", "0.0.0.0", "--port", "8080"]
"""
    )


def test_vetiver_write_app_rspm():
    file = "Dockerfile"
    vetiver.write_docker(rspm_env=True, path="./")
    contents = open(file).read()
    os.remove(file)
    py_version = str(sys.version_info.major) + "." + str(sys.version_info.minor)

    assert (
        contents
        == f"""# # Generated by the vetiver package; edit with care
# start with python base image
FROM python:{py_version}

# create directory in container for vetiver files
WORKDIR /vetiver

# copy  and install requirements
COPY vetiver_requirements.txt /vetiver/requirements.txt

#
RUN pip config set global.index-url https://colorado.rstudio.com/rspm/pypi/latest/simple
RUN pip install --no-cache-dir --upgrade -r /vetiver/requirements.txt

# copy app file
COPY app.py /vetiver/app/app.py

# expose port
EXPOSE 8080

# run vetiver API
CMD ["uvicorn", "app.app:api", "--host", "0.0.0.0", "--port", "8080"]
"""
    )
