Serialize a DisclosureOrigin to/from dict for task (#17276)

Author: di

tests/unit/integration/secrets/conftest.py

@@ -0,0 +1,25 @@
+# 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.
+
+import pytest
+
+from warehouse.integrations.secrets import utils
+
+
+@pytest.fixture
+def someorigin():
+    return utils.DisclosureOrigin(
+        name="SomeOrigin",
+        key_id_header="SOME_KEY_ID_HEADER",
+        signature_header="SOME_SIGNATURE_HEADER",
+        verification_url="https://some.verification.url",
+    )

tests/unit/integration/secrets/test_tasks.py

@@ -15,24 +15,23 @@
 from warehouse.integrations.secrets import tasks, utils
 
 
-def test_analyze_disclosure_task(monkeypatch):
+def test_analyze_disclosure_task(monkeypatch, someorigin):
     analyze_disclosure = pretend.call_recorder(lambda *a, **k: None)
     monkeypatch.setattr(utils, "analyze_disclosure", analyze_disclosure)
 
     request = pretend.stub()
     disclosure_record = pretend.stub()
-    origin = pretend.stub()
 
     tasks.analyze_disclosure_task(
         request=request,
         disclosure_record=disclosure_record,
-        origin=origin,
+        origin=someorigin.to_dict(),
     )
 
     assert analyze_disclosure.calls == [
         pretend.call(
             request=request,
             disclosure_record=disclosure_record,
-            origin=origin,
+            origin=someorigin,
         )
     ]

tests/unit/integration/secrets/test_utils.py

@@ -23,14 +23,31 @@
 from warehouse.integrations.secrets import tasks, utils
 
 
-@pytest.fixture
-def someorigin():
-    return utils.DisclosureOrigin(
-        name="SomeOrigin",
+def test_disclosure_origin_serialization(someorigin):
+    assert (
+        someorigin.to_dict()
+        == utils.DisclosureOrigin.from_dict(someorigin.to_dict()).to_dict()
+        == {
+            "name": "SomeOrigin",
+            "key_id_header": "SOME_KEY_ID_HEADER",
+            "signature_header": "SOME_SIGNATURE_HEADER",
+            "verification_url": "https://some.verification.url",
+            "api_token": None,
+        }
+    )
+
+
+def test_disclosure_origin_equivalence(someorigin):
+    assert someorigin == someorigin
+    someotherorigin = utils.DisclosureOrigin(
+        name="SomeOtherOrigin",
         key_id_header="SOME_KEY_ID_HEADER",
         signature_header="SOME_SIGNATURE_HEADER",
         verification_url="https://some.verification.url",
+        api_token=None,
     )
+    assert someorigin != someotherorigin
+    assert someorigin != "wu-tang"
 
 
 def test_token_leak_matcher_extract():
@@ -726,7 +743,7 @@ def test_analyze_disclosures_wrong_type(metrics, someorigin):
     assert exc.value.reason == "format"
 
 
-def test_analyze_disclosures_raise(metrics, monkeypatch):
+def test_analyze_disclosures_raise(metrics, monkeypatch, someorigin):
     task = pretend.stub(delay=pretend.call_recorder(lambda *a, **k: None))
     request = pretend.stub(task=lambda x: task)
 
@@ -735,12 +752,12 @@ def test_analyze_disclosures_raise(metrics, monkeypatch):
     utils.analyze_disclosures(
         request=request,
         disclosure_records=[1, 2, 3],
-        origin="yay",
+        origin=someorigin,
         metrics=metrics,
     )
 
     assert task.delay.calls == [
-        pretend.call(disclosure_record=1, origin="yay"),
-        pretend.call(disclosure_record=2, origin="yay"),
-        pretend.call(disclosure_record=3, origin="yay"),
+        pretend.call(disclosure_record=1, origin=someorigin.to_dict()),
+        pretend.call(disclosure_record=2, origin=someorigin.to_dict()),
+        pretend.call(disclosure_record=3, origin=someorigin.to_dict()),
     ]

warehouse/integrations/secrets/tasks.py

@@ -17,6 +17,7 @@
 
 @tasks.task(ignore_result=True, acks_late=True)
 def analyze_disclosure_task(request, disclosure_record, origin):
+    origin = utils.DisclosureOrigin.from_dict(origin)
     utils.analyze_disclosure(
         request=request,
         disclosure_record=disclosure_record,

warehouse/integrations/secrets/utils.py

@@ -50,6 +50,31 @@ def headers(self):
         """Set of all headers that must be present"""
         return {self.key_id_header, self.signature_header}
 
+    def to_dict(self):
+        return {
+            "name": self.name,
+            "key_id_header": self.key_id_header,
+            "signature_header": self.signature_header,
+            "verification_url": self.verification_url,
+            "api_token": self.api_token,
+        }
+
+    @classmethod
+    def from_dict(cls, data):
+        return cls(**data)
+
+    def __eq__(self, other):
+        if not isinstance(other, DisclosureOrigin):
+            return False
+
+        return (
+            self.name == other.name
+            and self.key_id_header == other.key_id_header
+            and self.signature_header == other.signature_header
+            and self.verification_url == other.verification_url
+            and self.api_token == other.api_token
+        )
+
 
 class ExtractionFailedError(Exception):
     pass
@@ -314,5 +339,5 @@ def analyze_disclosures(request, disclosure_records, origin, metrics):
 
     for disclosure_record in disclosure_records:
         request.task(tasks.analyze_disclosure_task).delay(
-            disclosure_record=disclosure_record, origin=origin
+            disclosure_record=disclosure_record, origin=origin.to_dict()
         )