feat: adapt python pb module to support pb stdout change (#438)

rudderlabs · Sep 3, 2024 · 2907172 · 2907172
1 parent 3509797
commit 2907172
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 2 deletions.
diff --git a/src/predictions/profiles_mlcorelib/wht/rudderPB.py b/src/predictions/profiles_mlcorelib/wht/rudderPB.py
@@ -84,7 +84,7 @@ def show_models(self, arg: dict) -> List[str]:
     def extract_json_from_stdout(self, stdout):
         start_index = stdout.find("printing models")
         if start_index == -1:
-            return None
+            start_index = 0
 
         # Find the index of the first '{' after the line
         start_index = stdout.find("{", start_index)

diff --git a/tests/unit/pythonWHT.py b/tests/unit/pythonWHT.py
@@ -171,7 +171,7 @@ def test_get_latest_seq_no(self):
 
 class TestGetInputs(unittest.TestCase):
     @patch("src.predictions.profiles_mlcorelib.wht.rudderPB.RudderPB.show_models")
-    def test_get_inputs(self, mock_rudderpb_show_models):
+    def test_get_inputs_from_stdout(self, mock_rudderpb_show_models):
         selector_sqls = [
             '''SELECT last_seen FROM "schema"."material_user_var_table_54ddc22a_383"''',
             """SELECT last_seen FROM `schema`.`material_user_var_table_54ddc22a_383`""",
@@ -273,3 +273,31 @@ def test_get_inputs(self, mock_rudderpb_show_models):
             },
         ]
         self.assertEqual(result, expected_result)
+
+    @patch("src.predictions.profiles_mlcorelib.wht.rudderPB.RudderPB.show_models")
+    def test_get_inputs_from_stderr(self, mock_rudderpb_show_models):
+        selector_sqls = [
+            '''SELECT last_seen FROM "schema"."material_user_var_table_54ddc22a_383"''',
+        ]
+        mock_rudderpb_show_models.return_value = """Some text before 
+            dummy entity.var
+            {
+                "project/user/all/last_seen": {
+                        "model_type": "entity_var_item"
+                }
+            }
+            Some text after"""
+        wht = PythonWHT("site_config_path", "project_folder")
+        result = wht.get_inputs(selector_sqls, False)
+        expected_result = [
+            {
+                "selector_sql": '''SELECT last_seen FROM "schema"."material_user_var_table_54ddc22a_383"''',
+                "table_name": "material_user_var_table_54ddc22a_383",
+                "model_name": "last_seen",
+                "column_name": "last_seen",
+                "model_ref": "user/all/last_seen",
+                "model_type": "entity_var_item",
+                "model_hash": None,
+            },
+        ]
+        self.assertEqual(result, expected_result)