fall back to the API when a datastream entity fetch times out

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);

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<String, String> 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