Fix deserialization of secret containung pydantic object

microsoft · Jun 15, 2021 · 1a18c50 · 1a18c50
1 parent 6b09456
commit 1a18c50
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 6 deletions.
diff --git a/src/api-service/tests/test_secrets.py b/src/api-service/tests/test_secrets.py
@@ -6,23 +6,26 @@
 import json
 import unittest
 
+from __app__.onefuzzlib.orm import hide_secrets
 from onefuzztypes.enums import OS, ContainerType
 from onefuzztypes.job_templates import (
     JobTemplate,
     JobTemplateIndex,
     JobTemplateNotification,
 )
 from onefuzztypes.models import (
+    GithubAuth,
+    GithubIssueTemplate,
     JobConfig,
     Notification,
     NotificationConfig,
+    NotificationTemplate,
     SecretAddress,
     SecretData,
     TeamsTemplate,
 )
 from onefuzztypes.primitives import Container
-
-from __app__.onefuzzlib.orm import hide_secrets
+from onefuzztypes.requests import NotificationCreate
 
 
 class TestSecret(unittest.TestCase):
@@ -91,3 +94,61 @@ def test_read_secret(self) -> None:
             self.assertIsInstance(notification.config.url.secret, SecretAddress)
         else:
             self.fail(f"Invalid config type {type(notification.config)}")
+
+    def test_read_secret2(self) -> None:
+        json_data = """
+            {
+                "notification_id": "b52b24d1-eec6-46c9-b06a-818a997da43c",
+                "container": "data",
+                "config" : {"url": {"secret": "http://test" }}
+            }
+            """
+        data = json.loads(json_data)
+        notification = Notification.parse_obj(data)
+        self.assertIsInstance(notification.config, TeamsTemplate)
+        if isinstance(notification.config, TeamsTemplate):
+            self.assertIsInstance(notification.config.url, SecretData)
+            self.assertIsInstance(notification.config.url.secret, str)
+        else:
+            self.fail(f"Invalid config type {type(notification.config)}")
+
+    def test_read_secret3(self) -> None:
+        json_data = """
+            {
+                "notification_id": "b52b24d1-eec6-46c9-b06a-818a997da43c",
+                "container": "data",
+                 "config": {
+                    "auth": {
+                        "secret": {
+                            "user": "INSERT_YOUR_USERNAME_HERE",
+                            "personal_access_token": "INSERT_YOUR_PERSONAL_ACCESS_TOKEN_HERE"
+                        }
+                    },
+                    "organization": "contoso",
+                    "repository": "sample-project",
+                    "title": "{{ report.executable }} - {{report.crash_site}}",
+                    "body": "",
+                    "unique_search": {
+                        "author": null,
+                        "state": null,
+                        "field_match": ["title"],
+                        "string": "{{ report.executable }}"
+                    },
+                    "assignees": [],
+                    "labels": ["bug", "{{ report.crash_type }}"],
+                    "on_duplicate": {
+                        "comment": "",
+                        "labels": ["{{ report.crash_type }}"],
+                        "reopen": true
+                    }
+                }
+            }
+            """
+        data = json.loads(json_data)
+        notification = Notification.parse_obj(data)
+        self.assertIsInstance(notification.config, GithubIssueTemplate)
+        if isinstance(notification.config, GithubIssueTemplate):
+            self.assertIsInstance(notification.config.auth, SecretData)
+            self.assertIsInstance(notification.config.auth.secret, GithubAuth)
+        else:
+            self.fail(f"Invalid config type {type(notification.config)}")
diff --git a/src/pytypes/onefuzztypes/models.py b/src/pytypes/onefuzztypes/models.py
@@ -4,7 +4,7 @@
 # Licensed under the MIT License.
 
 from datetime import datetime
-from typing import Any, Dict, Generic, List, Optional, Tuple, TypeVar, Union
+from typing import Any, Dict, Generic, List, Optional, Tuple, TypeVar, Union, get_args
 from uuid import UUID, uuid4
 
 from pydantic import BaseModel, Field, root_validator, validator
@@ -61,9 +61,9 @@ class SecretData(Generic[T]):
     secret: Union[T, SecretAddress]
 
     def __init__(self, secret: Union[T, SecretAddress]):
-        if isinstance(secret, dict):
+        try:
             self.secret = SecretAddress.parse_obj(secret)
-        else:
+        except Exception:
             self.secret = secret
 
     def __str__(self) -> str:
@@ -519,7 +519,7 @@ def validate_auth(cls, v: Any) -> SecretData:
             try:
                 return SecretData(GithubAuth.parse_obj(v))
             except Exception:
-                return SecretData(secret=v["secret"])
+                return SecretData(GithubAuth.parse_obj(v["secret"]))
         else:
             raise TypeError(f"invalid datatype {type(v)}")