Small annotations fix (#527)

OpenMined · Oct 16, 2024 · 4cf6a56 · 4cf6a56
1 parent 23e9d64
commit 4cf6a56
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 3 deletions.
diff --git 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
@@ -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