Pushing Environment externally to DataSource using ZIO

TestDSImplUsingZIO.scala

import zio.{DefaultRuntime, Ref, Task, UIO, ZIO}

type FieldSchema = String
type FieldData = String
type GenerateTimestamp = Long
type GenerateUuid[T] = T

type UserID=String
type AccountID = String

trait DataSourceData

case class UserProfile(userId:UserID, name:String)
  extends DataSourceData

case class AccountInfo(accountID:AccountID, accountName:String)
  extends DataSourceData

trait UserAPI {
  val user: UserAPI.Service
}
object UserAPI {
  // The User module contains the User service:
  trait Service {
    def getUser(id: String): Task[UserProfile]
  }
}

trait AccountAPI {
  val account: AccountAPI.Service
}
object AccountAPI {
  trait Service {
    def getAccount(id: String): Task[AccountInfo]
  }
}

trait DataSourceType
case object UserDataSource extends DataSourceType

type APIContext = UserAPI with AccountAPI

trait DataSourceTypeBehaviour[T <: DataSourceType]{
  def getData(data:T, requiredInput:Map[FieldSchema, FieldData])
             (implicit timestamp: GenerateTimestamp,
              uuid: GenerateUuid[String]):ZIO[APIContext, Throwable, Seq[DataSourceData]]
}

object UserDataSourceBehaviour extends DataSourceTypeBehaviour[UserDataSource.type]{
  def getData(data:UserDataSource.type, requiredInput:Map[FieldSchema, FieldData])
             (implicit timestamp: GenerateTimestamp, uuid: GenerateUuid[String])
  :ZIO[APIContext, Throwable, Seq[DataSourceData]] = {
    ZIO.accessM[APIContext] { apis =>
      import apis.user._
      getUser("abc").map(x => Seq(x))
    }
  }
}


object TestEnv{
  //Base trait for all the operations that are to be recorded when running test
  trait Operation

  //The TestState
  final case class ExecutedInstructions(ops: List[Operation])

  // A concurrent-safe test database service, which uses a `Ref` to keep track
  // of changes to the instruction set:
  class UserAPITestService(ref: Ref[ExecutedInstructions]) extends UserAPI.Service {
    private var map: Map[UserID, UserProfile] = Map("abc" -> UserProfile("abc", "testName"))

    override def getUser(id: String):Task[UserProfile] =
      ref.update(x => x.copy(ops = x.ops :+ UserAPITestService.UserLookedUp(id))).map(x => map(id))
  }
  object UserAPITestService {
    // database operations performed against the database:
    trait UserOperation extends Operation
    case class UserLookedUp(id:UserID) extends UserOperation
  }

  // A concurrent-safe test database service, which uses a `Ref` to keep track
  // of changes to the instruction set:
  class AccountAPITestService(ref: Ref[ExecutedInstructions]) extends AccountAPI.Service {
    private var map: Map[AccountID, AccountInfo] = Map("123" -> AccountInfo("abc", "Account123"))

    override def getAccount(id: AccountID):Task[AccountInfo] =
      ref.update(x => x.copy(ops = x.ops :+ AccountAPITestService.AccountLookedUp(id))).map(x => map(id))
  }
  object AccountAPITestService {
    // database operations performed against the database:
    trait AccountOperation extends Operation
    case class AccountLookedUp(id:AccountID) extends AccountOperation
  }


  // A helper function to run a test scenario, and extract out test data.
  // This function can be used many times across many unit tests.
  def testScenario[E, A](initialState: ExecutedInstructions)
                        (eff: ZIO[APIContext, E, A]): UIO[(Either[E, A], ExecutedInstructions)] = {
    for {
      instructionRef <- Ref.make(initialState)
      // Construct a new environment for the effect being tested:
      env = new UserAPI with AccountAPI {
        override val user = new UserAPITestService(instructionRef)
        override val account = new AccountAPITestService(instructionRef)
      }
      either <- eff.provide(env).either
      instructionState <- instructionRef.get
    } yield (either, instructionState)
  }
}

implicit val timeStamp:GenerateTimestamp = 123456L
implicit val uuidGen: GenerateUuid[String] = "testuuid"

val businessLogicResult = UserDataSourceBehaviour.getData(UserDataSource, Map.empty)
val v = TestEnv.testScenario(TestEnv.ExecutedInstructions(Nil))(businessLogicResult)
val runtime = new DefaultRuntime {}
runtime.unsafeRun(v)