From 301aa788d01ac65572a8ffa10e96b7e792cd87a7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Edgar=20Rami=CC=81rez=20Mondrago=CC=81n?=
 <edgarrm358@gmail.com>
Date: Thu, 11 May 2023 17:49:54 -0600
Subject: [PATCH] fix: Force stream selection in tests

---
 singer_sdk/streams/core.py                    | 11 +++++++++-
 singer_sdk/tap_base.py                        |  4 ++++
 tests/samples/test_tap_countries.py           | 22 +++++++++++++++++++
 .../countries_write_schemas                   |  2 ++
 4 files changed, 38 insertions(+), 1 deletion(-)
 create mode 100644 tests/snapshots/countries_write_schemas/countries_write_schemas

diff --git a/singer_sdk/streams/core.py b/singer_sdk/streams/core.py
index 31e16dc8b9..a414123ddc 100644
--- a/singer_sdk/streams/core.py
+++ b/singer_sdk/streams/core.py
@@ -269,7 +269,6 @@ def get_starting_timestamp(self, context: dict | None) -> datetime.datetime | No
 
         return t.cast(datetime.datetime, pendulum.parse(value))
 
-    @final
     @property
     def selected(self) -> bool:
         """Check if stream is selected.
@@ -279,6 +278,16 @@ def selected(self) -> bool:
         """
         return self.mask.get((), True)
 
+    @selected.setter
+    def selected(self, value: bool | None) -> None:
+        """Set stream selection.
+
+        Args:
+            value: True if the stream is selected.
+        """
+        self.metadata.root.selected = value
+        self._mask = self.metadata.resolve_selection()
+
     @final
     @property
     def has_selected_descendents(self) -> bool:
diff --git a/singer_sdk/tap_base.py b/singer_sdk/tap_base.py
index e6d999a021..a2d82ef988 100644
--- a/singer_sdk/tap_base.py
+++ b/singer_sdk/tap_base.py
@@ -246,6 +246,9 @@ def run_sync_dry_run(
             # Initialize streams' record limits before beginning the sync test.
             stream.ABORT_AT_RECORD_COUNT = dry_run_record_limit
 
+            # Force selection of streams.
+            stream.selected = True
+
         for stream in streams:
             if stream.parent_stream_type:
                 self.logger.debug(
@@ -267,6 +270,7 @@ def run_sync_dry_run(
     def write_schemas(self) -> None:
         """Write a SCHEMA message for all known streams to STDOUT."""
         for stream in self.streams.values():
+            stream.selected = True
             stream._write_schema_message()
 
     # Stream detection:
diff --git a/tests/samples/test_tap_countries.py b/tests/samples/test_tap_countries.py
index 5e7a5cb986..7fc48b454e 100644
--- a/tests/samples/test_tap_countries.py
+++ b/tests/samples/test_tap_countries.py
@@ -9,6 +9,9 @@
 import typing as t
 from contextlib import redirect_stdout
 
+import pytest
+from click.testing import CliRunner
+
 from samples.sample_tap_countries.countries_tap import SampleTapCountries
 from singer_sdk.helpers._catalog import (
     get_selected_schema,
@@ -17,6 +20,11 @@
 from singer_sdk.testing import get_tap_test_class
 from singer_sdk.testing.config import SuiteConfig
 
+if t.TYPE_CHECKING:
+    from pathlib import Path
+
+    from pytest_snapshot.plugin import Snapshot
+
 SAMPLE_CONFIG = {}
 SAMPLE_CONFIG_BAD = {"not": "correct"}
 
@@ -136,3 +144,17 @@ def tally_messages(messages: list) -> t.Counter:
     assert counter["BATCH", "countries"] == 1
 
     assert counter[("STATE",)] == 2
+
+
+@pytest.mark.snapshot()
+def test_write_schema(
+    snapshot: Snapshot,
+    snapshot_dir: Path,
+):
+    snapshot.snapshot_dir = snapshot_dir.joinpath("countries_write_schemas")
+
+    runner = CliRunner(mix_stderr=False)
+    result = runner.invoke(SampleTapCountries.cli, ["--test", "schema"])
+
+    snapshot_name = "countries_write_schemas"
+    snapshot.assert_match(result.stdout, snapshot_name)
diff --git a/tests/snapshots/countries_write_schemas/countries_write_schemas b/tests/snapshots/countries_write_schemas/countries_write_schemas
new file mode 100644
index 0000000000..b0808ce23c
--- /dev/null
+++ b/tests/snapshots/countries_write_schemas/countries_write_schemas
@@ -0,0 +1,2 @@
+{"type": "SCHEMA", "stream": "continents", "schema": {"properties": {"code": {"type": ["null", "string"]}, "name": {"type": ["null", "string"]}}, "type": "object"}, "key_properties": ["code"]}
+{"type": "SCHEMA", "stream": "countries", "schema": {"properties": {"code": {"type": ["string", "null"]}, "name": {"type": ["string", "null"]}, "native": {"type": ["string", "null"]}, "phone": {"type": ["string", "null"]}, "capital": {"type": ["string", "null"]}, "currency": {"type": ["string", "null"]}, "emoji": {"type": ["string", "null"]}, "continent": {"properties": {"code": {"type": ["string", "null"]}, "name": {"type": ["string", "null"]}}, "type": ["object", "null"]}, "languages": {"items": {"properties": {"code": {"type": ["string", "null"]}, "name": {"type": ["string", "null"]}}, "type": "object"}, "type": ["array", "null"]}}, "type": "object"}, "key_properties": ["code"]}