test: Fix tests

Author: erezrokah

cloudquery/sdk/internal/memdb/memdb.py

@@ -7,6 +7,7 @@
 from typing import List, Generator, Dict, Any
 import pyarrow as pa
 from cloudquery.sdk.schema.table import Table
+from cloudquery.sdk.schema.arrow import METADATA_TABLE_NAME
 from cloudquery.sdk.types import JSONType
 from dataclasses import dataclass, field
 
@@ -155,7 +156,7 @@ def __init__(self) -> None:
             table1.name: table1,
             table2.name: table2,
         }
-        self._db: Dict[str, pa.RecordBatch] = {}
+        self._db: List[pa.RecordBatch] = []
         self._client = Client()
 
     def set_logger(self, logger) -> None:
@@ -202,18 +203,15 @@ def sync(
     def write(self, writer: Generator[message.WriteMessage, None, None]) -> None:
         for msg in writer:
             if isinstance(msg, message.WriteMigrateTableMessage):
-                if msg.table.name not in self._db:
-                    self._db[msg.table.name] = msg.table
-                    self._tables[msg.table.name] = msg.table
+                pass
             elif isinstance(msg, message.WriteInsertMessage):
-                table = schema.Table.from_arrow_schema(msg.record.schema)
-                self._db[table.name] = msg.record
+                self._db.append(msg.record)
             else:
                 raise NotImplementedError(f"Unknown message type {type(msg)}")
 
     def read(self, table: Table) -> Generator[message.ReadMessage, None, None]:
-        for table, record in self._db.items():
-            recordMetadata = record.schema.metadata.get(schema.MetadataTableName)
+        for record in self._db:
+            recordMetadata = record.schema.metadata.get(METADATA_TABLE_NAME).decode("utf-8")
             if recordMetadata == table.name:
                 yield message.ReadMessage(record)
 

tests/internal/memdb/memdb.py

@@ -1,12 +1,25 @@
 from cloudquery.sdk.internal import memdb
 from cloudquery.sdk.internal.servers.plugin_v3 import plugin
 from cloudquery.sdk.plugin import SyncOptions
+from cloudquery.sdk.message import SyncMigrateTableMessage, SyncInsertMessage
+import structlog
 
 
 def test_memdb():
     p = memdb.MemDB()
+    p.set_logger(structlog.get_logger())
     p.init(plugin.sanitize_spec(b"null"))
     msgs = []
-    for msg in p.sync(SyncOptions(tables=["*"])):
+    for msg in p.sync(SyncOptions(tables=["*"],skip_tables=[])):
         msgs.append(msg)
-    assert len(msgs) == 0
+    assert len(msgs) == 18
+
+    assert isinstance(msgs[0], SyncMigrateTableMessage)
+    assert isinstance(msgs[1], SyncMigrateTableMessage)
+    assert isinstance(msgs[2], SyncMigrateTableMessage)
+
+    # other messages should be inserts
+    for msg in msgs[3:]:
+        assert isinstance(msg, SyncInsertMessage)
+
+

tests/serve/plugin.py

@@ -14,15 +14,9 @@
 from cloudquery.sdk.types.uuid import UUIDType
 
 test_table = Table(
-    "test",
-    [
-        Column("id", pa.int64()),
-        Column("name", pa.string()),
-        Column("json", JSONType()),
-        Column("uuid", UUIDType()),
-    ],
-)
-
+        "test_table",
+        [Column("id", pa.int64()), Column("name", pa.string()), Column("json", JSONType()), Column("uuid", UUIDType())],
+    )
 
 def test_plugin_serve():
     p = MemDB()
@@ -72,17 +66,29 @@ def writer_iterator():
 
             stub.Write(writer_iterator())
 
-            response = stub.GetTables(plugin_pb2.GetTables.Request(tables=["*"]))
+            response = stub.GetTables(plugin_pb2.GetTables.Request(tables=["*"],skip_tables=[]))
             schemas = arrow.new_schemas_from_bytes(response.tables)
-            assert len(schemas) == 4
+            assert len(schemas) == 3
 
-            response = stub.Sync(plugin_pb2.Sync.Request(tables=["*"]))
+            response = stub.Sync(plugin_pb2.Sync.Request(tables=["*"],skip_tables=[]))
+            total_migrate_tables = 0
             total_records = 0
+            total_errors = 0
             for msg in response:
-                if msg.insert is not None:
+                message_type = msg.WhichOneof("message")
+                if message_type == "insert":
                     rec = arrow.new_record_from_bytes(msg.insert.record)
+                    assert rec.num_rows > 0
                     total_records += 1
-            assert total_records == 1
+                elif message_type == "migrate_table":
+                    total_migrate_tables += 1
+                elif message_type == "error":
+                    total_errors += 1
+                else:
+                    raise NotImplementedError(f"Unknown message type {type(msg)}")
+            assert total_migrate_tables == 3
+            assert total_records == 15
+            assert total_errors == 0
     finally:
         cmd.stop()
         pool.shutdown()
@@ -122,8 +128,7 @@ def test_plugin_read():
         ],
         schema=test_table.to_arrow_schema(),
     )
-    p._db["test_1"] = sample_record_1
-    p._db["test_2"] = sample_record_2
+    p._db = [sample_record_1, sample_record_2]
 
     cmd = serve.PluginCommand(p)
     port = random.randint(5000, 50000)
@@ -191,7 +196,7 @@ def test_plugin_package():
                     },
                     {
                         "name": "id",
-                        "type": "string",
+                        "type": "int64",
                         "description": "",
                         "incremental_key": True,
                         "primary_key": True,
@@ -247,7 +252,7 @@ def test_plugin_package():
                     },
                     {
                         "name": "id",
-                        "type": "string",
+                        "type": "int64",
                         "description": "",
                         "incremental_key": False,
                         "primary_key": False,