From 5c51398f69e9298e9a4b3293f5e43e433ac834ed Mon Sep 17 00:00:00 2001
From: Josh Callender <1569818+saponifi3d@users.noreply.github.com>
Date: Fri, 13 Dec 2024 17:49:52 -0800
Subject: [PATCH] add some more tests to make sure everything is working as
 expected

---
 .../models/test_data_source.py                | 23 +++++++++++++++++++
 .../processors/test_data_sources.py           |  6 +++++
 2 files changed, 29 insertions(+)
 create mode 100644 tests/sentry/workflow_engine/models/test_data_source.py

diff --git a/tests/sentry/workflow_engine/models/test_data_source.py b/tests/sentry/workflow_engine/models/test_data_source.py
new file mode 100644
index 00000000000000..2c9d9a9e86b6d0
--- /dev/null
+++ b/tests/sentry/workflow_engine/models/test_data_source.py
@@ -0,0 +1,23 @@
+from unittest import mock
+
+import pytest
+
+from sentry.workflow_engine.registry import data_source_type_registry
+from tests.sentry.workflow_engine.test_base import BaseWorkflowTest
+
+
+class DataSourceTest(BaseWorkflowTest):
+    def test_inavlid_data_source_type(self):
+        with pytest.raises(ValueError):
+            self.create_data_source(type="invalid_type")
+
+    def test_data_source_valid_type(self):
+        test_type = "test"
+        type_mock = mock.Mock()
+
+        data_source_type_registry.register(test_type)(type_mock)
+        assert data_source_type_registry.get(test_type) == type_mock
+
+        data_source = self.create_data_source(type=test_type)
+
+        assert data_source is not None
diff --git a/tests/sentry/workflow_engine/processors/test_data_sources.py b/tests/sentry/workflow_engine/processors/test_data_sources.py
index 2e655b67007085..56b44e7b2798af 100644
--- a/tests/sentry/workflow_engine/processors/test_data_sources.py
+++ b/tests/sentry/workflow_engine/processors/test_data_sources.py
@@ -1,7 +1,13 @@
+from unittest import mock
+
 from sentry.snuba.models import SnubaQuery
 from sentry.testutils.cases import TestCase
 from sentry.workflow_engine.models import DataPacket
 from sentry.workflow_engine.processors import process_data_sources
+from sentry.workflow_engine.registry import data_source_type_registry
+
+# Setup a mock data source for the tests
+data_source_type_registry.register("test")(mock.Mock)
 
 
 class TestProcessDataSources(TestCase):