/*!
 * Copyright 2021 by Contributors
 */
#include <gtest/gtest.h>
#include <xgboost/host_device_vector.h>
#include "../helpers.h"
#include "../../../src/data/array_interface.h"

namespace xgboost {

__global__ void SleepForTest(uint64_t *out, uint64_t duration) {
  auto start = clock64();
  auto t = 0;
  while (t < duration) {
    t = clock64() - start;
  }
  out[0] = t;
}

TEST(ArrayInterface, Stream) {
  size_t constexpr kRows = 10, kCols = 10;
  HostDeviceVector<float> storage;
  auto arr_str = RandomDataGenerator{kRows, kCols, 0}.GenerateArrayInterface(&storage);

  cudaStream_t stream;
  cudaStreamCreate(&stream);

  auto j_arr =Json::Load(StringView{arr_str});
  j_arr["stream"] = Integer(reinterpret_cast<int64_t>(stream));
  Json::Dump(j_arr, &arr_str);

  dh::caching_device_vector<uint64_t> out(1, 0);
  uint64_t dur = 1e9;
  dh::LaunchKernel{1, 1, 0, stream}(SleepForTest, out.data().get(), dur);
  ArrayInterface arr(arr_str);

  auto t = out[0];
  CHECK_GE(t, dur);

  cudaStreamDestroy(stream);
}
}  // namespace xgboost