src/algorithms/privacy/poisson_test.cc

#include "src/algorithms/privacy/poisson.h"

#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include "src/algorithms/random/test_secure_random.h"

using testing::DoubleNear;
using testing::UnorderedElementsAreArray;

namespace cobalt {
class PoissonNoiseTest : public ::testing::Test {
 protected:
  void SetUp() override { gen_ = std::make_unique<TestSecureRandomNumberGenerator>(0); }

  TestSecureRandomNumberGenerator* GetGenerator() { return gen_.get(); }

 private:
  std::unique_ptr<TestSecureRandomNumberGenerator> gen_;
};

TEST_F(PoissonNoiseTest, LambdaIsZero) {
  std::vector<uint64_t> indices = {2, 3, 5, 8, 10};
  int64_t max_index = 10;
  PoissonParameter lambda(0.0);
  auto noised = ApplyPoissonNoise(indices, max_index, lambda, GetGenerator());
  EXPECT_THAT(noised, UnorderedElementsAreArray(indices));
}

// Check that the total number of indices changed over 1000 trials is within 3
// standard deviations of the expected number.
TEST_F(PoissonNoiseTest, NumIndicesChanged) {
  std::vector<uint64_t> indices;
  uint64_t num_trials = 1000;
  size_t total_indices_changed = 0;

  int64_t max_index = 99;
  PoissonParameter lambda(0.1);

  for (uint64_t i = 0; i < num_trials; i++) {
    total_indices_changed += ApplyPoissonNoise(indices, max_index, lambda, GetGenerator()).size();
  }

  double stddev = std::sqrt(static_cast<double>(num_trials * (1 + max_index)) * lambda.get());
  double interval_radius = 3 * stddev;
  double expected_indices_changed =
      static_cast<double>(num_trials * (1 + max_index)) * lambda.get();
  EXPECT_THAT(static_cast<double>(total_indices_changed),
              DoubleNear(expected_indices_changed, interval_radius));
}

TEST_F(PoissonNoiseTest, DebiasPoissonCountPIsZero) {
  PoissonParameter lambda(0.0);
  double raw_count = 500.0;
  uint64_t num_contributions = 1000;

  EXPECT_THAT(DebiasPoissonCount(raw_count, num_contributions, lambda),
              DoubleNear(static_cast<double>(raw_count), 0.001));
}

TEST_F(PoissonNoiseTest, DebiasPoissonCountNegativeEstimate) {
  PoissonParameter lambda(0.25);
  double raw_count = 500.0;
  uint64_t num_contributions = 10000;

  EXPECT_THAT(DebiasPoissonCount(raw_count, num_contributions, lambda), DoubleNear(-2000, 0.001));
}

TEST_F(PoissonNoiseTest, DebiasPoissonCount) {
  PoissonParameter lambda(0.1);
  double raw_count = 2000.0;
  uint64_t num_contributions = 10000;

  EXPECT_THAT(DebiasPoissonCount(raw_count, num_contributions, lambda), DoubleNear(1000, 0.001));
}

}  // namespace cobalt