#!/bin/bash

HOME_DIR=/test
BIN_DIR=${HOME_DIR}/bin
LOG_DIR=${HOME_DIR}/logs

TORCHDATA_BRANCH_DEFAULT="main"
TORCHDATA_REPO="https://github.com/pytorch/data.git"
TORCHDATA_TEST_SCRIPT="data/test/test_remote_io.py"

echo "Fetch Torchdata S3 IO test script."
TORCHDATA_RELEASE_TAG=$(python -c "import torchdata as td; print(td.__version__.split('+')[0])")
git clone $TORCHDATA_REPO \
    && pushd data \
    && git checkout tags/v${TORCHDATA_RELEASE_TAG} -b ${TORCHDATA_RELEASE_TAG} \
    && popd

echo "Installing pre-requisites."
pip install -r data/test/requirements.txt

echo "Running Torchdata S3 IO datapipe tests."
pytest data/test/test_remote_io.py

if [ "$?" -eq "0" ]
then
    echo "Torchdata S3 IO tests succeeded."
else
    echo "Torchdata S3 IO tests failed at ${TORCHDATA_REPO} ${TORCHDATA_BRANCH} ${TORCHDATA_TEST_SCRIPT}"
    exit 1
fi

exit 0