From eda6b8e36eef0e868d17516cfa465ca571e56af0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Vinot?= <benoit.vinot@roseautechnologies.com>
Date: Tue, 10 Jan 2023 09:53:21 +0100
Subject: [PATCH] Add tests for optimise

---
 mapbox_vector_tile/optimise.py |  8 ++---
 pyproject.toml                 |  4 +--
 tests/test_optimise.py         | 65 ++++++++++++++++++++++++++++++++++
 3 files changed, 70 insertions(+), 7 deletions(-)
 create mode 100644 tests/test_optimise.py

diff --git a/mapbox_vector_tile/optimise.py b/mapbox_vector_tile/optimise.py
index 05753d1..57d23a5 100644
--- a/mapbox_vector_tile/optimise.py
+++ b/mapbox_vector_tile/optimise.py
@@ -1,7 +1,7 @@
 from collections import namedtuple
 
-from mapbox_vector_tile.Mapbox.vector_tile_pb2 import tile
-from mapbox_vector_tile.utils import CMD_LINE_TO, CMD_MOVE_TO, zig_zag_decode, zig_zag_encode
+from mapbox_vector_tile.Mapbox import vector_tile_pb2 as vector_tile
+from mapbox_vector_tile.utils import CMD_LINE_TO, CMD_MOVE_TO, LINESTRING, zig_zag_decode, zig_zag_encode
 
 
 class StringTableOptimiser:
@@ -208,7 +208,7 @@ def optimise_tile(tile_bytes):
     multilinestrings to save a few bytes.
     """
 
-    t = tile()
+    t = vector_tile.tile()
     t.ParseFromString(tile_bytes)
 
     for layer in t.layers:
@@ -216,7 +216,7 @@ def optimise_tile(tile_bytes):
 
         for feature in layer.features:
             # (multi)linestrings only
-            if feature.type == 2:
+            if feature.type == LINESTRING:
                 optimise_multilinestring(feature.geometry)
 
             sto.add_tags(feature.tags)
diff --git a/pyproject.toml b/pyproject.toml
index 58366a4..cba7fe1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -51,9 +51,7 @@ line-length = 120
 fast = true
 extend-exclude = '''
 (
-  /(
-    mapbox_vector_tile\/Mapbox\/vector_tile_pb2.py
-  )/
+    .*_pb2.py
 )
 '''
 target-version = ["py37", "py38", "py39", "py310", "py311"]
diff --git a/tests/test_optimise.py b/tests/test_optimise.py
new file mode 100644
index 0000000..bbf0974
--- /dev/null
+++ b/tests/test_optimise.py
@@ -0,0 +1,65 @@
+import unittest
+
+import mapbox_vector_tile
+from mapbox_vector_tile.Mapbox import vector_tile_pb2 as vector_tile
+from mapbox_vector_tile.optimise import optimise_tile
+
+
+class BaseTestCase(unittest.TestCase):
+    def test_string_table_optimizer(self):
+        tile_data = mapbox_vector_tile.encode(
+            {
+                "name": "water",
+                "features": [
+                    {"geometry": "POLYGON ((0 0, 0 1, 1 1, 1 0, 0 0))", "properties": {"uid": 123, "foo": "bar"}},
+                    {"geometry": "POLYGON ((0 0, 0 1, 1 1, 1 0, 0 0))", "properties": {"uid": 124, "foo": "bar"}},
+                    {
+                        "geometry": "POLYGON ((0 0, 0 1, 1 1, 1 0, 0 0))",
+                        "properties": {"uid": 125, "cat": "flew", "foo": "bar"},
+                    },
+                ],
+            }
+        )
+        # Check the tags in the encoded tile
+        tile = vector_tile.tile()
+        tile.ParseFromString(tile_data)
+        tile_layer = tile.layers[0]
+        expected_tile_tags = ([0, 0, 1, 1], [0, 2, 1, 1], [0, 3, 2, 4, 1, 1])
+        for i, f in enumerate(tile_layer.features):
+            self.assertEqual(f.tags, expected_tile_tags[i])
+
+        # Optimise it
+        result = optimise_tile(tile_data)
+
+        # Check that the tags have been modified
+        result_tile = vector_tile.tile()
+        result_tile.ParseFromString(result)
+        result_layer = result_tile.layers[0]
+        expected_results = ([1, 4, 0, 0], [1, 3, 0, 0], [1, 2, 2, 1, 0, 0])
+        for i, f in enumerate(result_layer.features):
+            self.assertEqual(f.tags, expected_results[i])
+
+    def test_optimise_linestring(self):
+        # No particular optimization as only a linestring
+        tile_data = mapbox_vector_tile.encode(
+            {"name": "water", "features": [{"geometry": "LINESTRING (0 0, 0 1, 1 1, 1 0, 0 0)", "properties": {}}]}
+        )
+        result = optimise_tile(tile_data)
+        self.assertEqual(tile_data, result)
+
+        # Multilinestring to optimize
+        tile_data = mapbox_vector_tile.encode(
+            {
+                "name": "water",
+                "features": [
+                    {"geometry": "MULTILINESTRING ((0 0, 0 1, 1 1), (2 2, 3 2), (1 1, 2 2))", "properties": {}}
+                ],
+            }
+        )
+        result = optimise_tile(tile_data)
+        decoded_result = mapbox_vector_tile.decode(result)
+        decoded_geometry = decoded_result["water"]["features"][0]["geometry"]
+        self.assertEqual(decoded_geometry["type"], "MultiLineString")
+        self.assertEqual(
+            decoded_geometry["coordinates"], [[[0, 0], [0, 1], [1, 1]], [[1, 1], [2, 2]], [[2, 2], [3, 2]]]
+        )