diff --git a/src/main/java/com/schematic/api/datastream/DataStreamClient.java b/src/main/java/com/schematic/api/datastream/DataStreamClient.java index fb848785..9c98e185 100644 --- a/src/main/java/com/schematic/api/datastream/DataStreamClient.java +++ b/src/main/java/com/schematic/api/datastream/DataStreamClient.java @@ -252,11 +252,19 @@ public RulesengineCheckFlagResult checkFlag( throw new DataStreamException("Datastream not connected and required entities not in cache"); } + // Evaluating without an entity the caller named would skip its targeted rules and report + // the result as authoritative, so a failed fetch throws and the caller falls back to the API. if (needsCompany && cachedCompany == null) { cachedCompany = getCompany(company); + if (cachedCompany == null) { + throw new DataStreamException("Company could not be fetched from datastream"); + } } if (needsUser && cachedUser == null) { cachedUser = getUser(user); + if (cachedUser == null) { + throw new DataStreamException("User could not be fetched from datastream"); + } } return evaluateFlag(flag, cachedCompany, cachedUser, preflight); diff --git a/src/test/java/com/schematic/api/datastream/DataStreamClientTest.java b/src/test/java/com/schematic/api/datastream/DataStreamClientTest.java index 982b78a1..5ae6c537 100644 --- a/src/test/java/com/schematic/api/datastream/DataStreamClientTest.java +++ b/src/test/java/com/schematic/api/datastream/DataStreamClientTest.java @@ -16,6 +16,7 @@ import com.schematic.api.types.RulesengineCompany; import com.schematic.api.types.RulesengineFlag; import com.schematic.api.types.RulesengineUser; +import java.lang.reflect.Field; import java.time.Duration; import java.util.Collections; import java.util.HashMap; @@ -327,6 +328,27 @@ void checkFlag_withCachedCompanyAndUser() { assertEquals("user-1", result.getUserId().orElse(null)); } + @Test + void checkFlag_throwsWhenCompanyFetchTimesOut() throws Exception { + DataStreamClient connected = new DataStreamClient( + DatastreamOptions.builder().build(), "test-key", "https://api.schematichq.com", logger); + DataStreamWebSocketClient ws = mock(DataStreamWebSocketClient.class); + when(ws.isReady()).thenReturn(true); + Field wsField = DataStreamClient.class.getDeclaredField("wsClient"); + wsField.setAccessible(true); + wsField.set(connected, ws); + + connected.handleMessage(buildResp( + EntityType.FLAG.getValue(), MessageType.FULL.getValue(), null, flagNode("company-flag", true))); + + Map companyKeys = Collections.singletonMap("customer_id", "never-arrives"); + assertThrows(DataStreamException.class, () -> connected.checkFlag("company-flag", companyKeys, null)); + verify(logger, times(1)).warn(anyString()); + + wsField.set(connected, null); + connected.close(); + } + @Test void handleMessage_partialCompany_noExisting_fallsBackToFullParse() { // Send a PARTIAL company message when nothing is cached for that entity_id