<commit_before>/*
 * Copyright 2016 Spotify AB.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package com.spotify.scio.util.random

import com.spotify.scio.testing.PipelineSpec
import org.apache.beam.sdk.options.PipelineOptions
import org.apache.beam.sdk.transforms.DoFn
import org.apache.beam.sdk.transforms.windowing.PaneInfo
import org.apache.beam.sdk.values.{PCollectionView, TupleTag}
import org.joda.time.Instant

import scala.collection.mutable.{Buffer => MBuffer}

class RandomSamplerTest extends PipelineSpec {

  import RandomSamplerUtils._

  def test[T](sampler: RandomSampler[T, _], xs: Seq[T]): Seq[T] = {
    sampler.startBundle(null)
    val buffer = MBuffer.empty[T]
    xs.foreach(x => sampler.processElement(newContext[T](sampler, x, buffer)))
    buffer
  }

  def test[K, V](sampler: RandomValueSampler[K, V, _], xs: Seq[(K, V)]): Seq[(K, V)] = {
    sampler.startBundle(null)
    val buffer = MBuffer.empty[(K, V)]
    xs.foreach(x => sampler.processElement(newContext[(K, V)](sampler, x, buffer)))
    buffer
  }

  private def newContext[T](f: DoFn[T, T], e: T, buffer: MBuffer[T]) = new f.ProcessContext {
    override def element(): T = e
    override def sideInput[U](view: PCollectionView[U]): U = ???
    override def timestamp(): Instant = ???
    override def pane(): PaneInfo = ???
    override def updateWatermark(watermark: Instant): Unit = ???
    override def getPipelineOptions: PipelineOptions = ???
    override def output(output: T): Unit = buffer.append(output)
    override def output[U](tag: TupleTag[U], output: U): Unit = ???
    override def outputWithTimestamp(output: T, timestamp: Instant): Unit = ???
    override def outputWithTimestamp[U](tag: TupleTag[U], output: U, timestamp: Instant): Unit = ???
  }

  def testSampler(withReplacement: Boolean,
                  expectedFraction: Double, actualFraction: Double): Double = {
    rngSeed.setSeed(fixedSeed)
    val expected = expectedSamples(withReplacement, expectedFraction)

    val sampler = if (withReplacement) {
      new PoissonSampler[Int](actualFraction)<commit_msg>fix RandomSamplerTest method visibility<commit_after>/*
 * Copyright 2016 Spotify AB.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package com.spotify.scio.util.random

import com.spotify.scio.testing.PipelineSpec
import org.apache.beam.sdk.options.PipelineOptions
import org.apache.beam.sdk.transforms.DoFn
import org.apache.beam.sdk.transforms.windowing.PaneInfo
import org.apache.beam.sdk.values.{PCollectionView, TupleTag}
import org.joda.time.Instant

import scala.collection.mutable.{Buffer => MBuffer}

class RandomSamplerTest extends PipelineSpec {

  import RandomSamplerUtils._

  private def test[T](sampler: RandomSampler[T, _], xs: Seq[T]): Seq[T] = {
    sampler.startBundle(null)
    val buffer = MBuffer.empty[T]
    xs.foreach(x => sampler.processElement(newContext[T](sampler, x, buffer)))
    buffer
  }

  private def test[K, V](sampler: RandomValueSampler[K, V, _], xs: Seq[(K, V)]): Seq[(K, V)] = {
    sampler.startBundle(null)
    val buffer = MBuffer.empty[(K, V)]
    xs.foreach(x => sampler.processElement(newContext[(K, V)](sampler, x, buffer)))
    buffer
  }

  private def newContext[T](f: DoFn[T, T], e: T, buffer: MBuffer[T]) = new f.ProcessContext {
    override def element(): T = e
    override def sideInput[U](view: PCollectionView[U]): U = ???
    override def timestamp(): Instant = ???
    override def pane(): PaneInfo = ???
    override def updateWatermark(watermark: Instant): Unit = ???
    override def getPipelineOptions: PipelineOptions = ???
    override def output(output: T): Unit = buffer.append(output)
    override def output[U](tag: TupleTag[U], output: U): Unit = ???
    override def outputWithTimestamp(output: T, timestamp: Instant): Unit = ???
    override def outputWithTimestamp[U](tag: TupleTag[U], output: U, timestamp: Instant): Unit = ???
  }

  def testSampler(withReplacement: Boolean,
                  expectedFraction: Double, actualFraction: Double): Double = {
    rngSeed.setSeed(fixedSeed)
    val expected = expectedSamples(withReplacement, expectedFraction)

    val sampler = if (withReplacement) {
      new PoissonSampler[Int](actualFraction)
{'fraction_non_alphanumeric': 0.12047274113610369, 'fraction_numerical': 0.0030499428135722455, 'mean_word_length': 4.2260956175298805, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5834434', 'n_tokens_mistral': 1637, 'n_tokens_neox': 1581, 'n_words': 462}