From 4cf6a562693299792100c023e963d7c0ca89566c Mon Sep 17 00:00:00 2001
From: Vadym Doroshenko <53558779+dvadym@users.noreply.github.com>
Date: Wed, 16 Oct 2024 11:45:12 +0200
Subject: [PATCH] Small annotations fix (#527)

---
 pipeline_dp/dataset_histograms/sum_histogram_computation.py | 2 +-
 tests/dp_engine_test.py                                     | 2 --
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/pipeline_dp/dataset_histograms/sum_histogram_computation.py b/pipeline_dp/dataset_histograms/sum_histogram_computation.py
index 63d788bf..baa43136 100644
--- a/pipeline_dp/dataset_histograms/sum_histogram_computation.py
+++ b/pipeline_dp/dataset_histograms/sum_histogram_computation.py
@@ -112,7 +112,7 @@ def _compute_frequency_histogram_per_key(
     def _map_to_frequency_bin(
         key_value: Tuple[int, float],
         bucket_generators: List[List[LowerUpperGenerator]]
-    ) -> Tuple[float, hist.FrequencyBin]:
+    ) -> Tuple[Tuple[int, float], hist.FrequencyBin]:
         # bucket_generator is a 1-element list with
         # a single element to be a list of LowerUpperGenerator.
         index, value = key_value
diff --git a/tests/dp_engine_test.py b/tests/dp_engine_test.py
index 150d6d73..6169ad9b 100644
--- a/tests/dp_engine_test.py
+++ b/tests/dp_engine_test.py
@@ -254,8 +254,6 @@ def test_calculate_private_contribution_filters_partitions(self):
             result,
             pipeline_dp.PrivateContributionBounds(max_partitions_contributed=1))
 
-    @unittest.skip(
-        "For some reason it fails on Beam. TODO(dvadym): Fix this test")
     def test_calculate_private_contribution_works_on_beam(self):
         with test_pipeline.TestPipeline() as p:
             # Arrange